check torch version inside functions

Signed-off-by: Xin Yao <[email protected]>
NVIDIA · Oct 23, 2024 · d162789 · d162789
1 parent 39a9fe6
commit d162789
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()