check torch version inside functions

Signed-off-by: Xin Yao <[email protected]>
NVIDIA · Nov 5, 2024 · 64a506c · 64a506c
1 parent d2f5be8
commit 64a506c
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/transformer_engine/pytorch/utils.py b/transformer_engine/pytorch/utils.py
@@ -313,5 +313,4 @@ def torch_get_autocast_gpu_dtype() -> torch.dtype:
     TORCH_MINOR = int(torch.__version__.split(".")[1])
     if TORCH_MAJOR == 2 and TORCH_MINOR >= 4:
         return torch.get_autocast_dtype("cuda")
-    else:
-        return torch.get_autocast_gpu_dtype()
+    return torch.get_autocast_gpu_dtype()