fix kwargs for torch.amp.autocast

Signed-off-by: Xin Yao <[email protected]>
NVIDIA · Oct 21, 2024 · 35d7a31 · 35d7a31
1 parent 0987833
commit 35d7a31
Showing 1 changed file with 8 additions and 2 deletions.
diff --git a/transformer_engine/pytorch/distributed.py b/transformer_engine/pytorch/distributed.py
@@ -259,13 +259,19 @@ def _get_active_autocast_contexts():
         gpu_autocast_enabled = torch.is_autocast_enabled("cuda")
         gpu_autocast_dtype = torch.get_autocast_dtype("cuda")
         gpu_autocast_ctx = torch.amp.autocast(
-            "cuda", gpu_autocast_enabled, gpu_autocast_dtype, autocast_cached
+            "cuda",
+            enabled=gpu_autocast_enabled,
+            dtype=gpu_autocast_dtype,
+            cache_enabled=autocast_cached,
         )
 
         cpu_autocast_enabled = torch.is_autocast_enabled("cpu")
         cpu_autocast_dtype = torch.get_autocast_dtype("cpu")
         cpu_autocast_ctx = torch.amp.autocast(
-            "cpu", cpu_autocast_enabled, cpu_autocast_dtype, autocast_cached
+            "cpu",
+            enabled=cpu_autocast_enabled,
+            dtype=cpu_autocast_dtype,
+            cache_enabled=autocast_cached,
         )
 
         return gpu_autocast_ctx, cpu_autocast_ctx