Update apex/transformer/utils.py

Co-authored-by: Masaki Kozuki <[email protected]>
NVIDIA · May 9, 2024 · bfeb06a · bfeb06a
1 parent 64aea56
commit bfeb06a
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/apex/transformer/utils.py b/apex/transformer/utils.py
@@ -8,7 +8,8 @@
 # The following 4 lines are for backward comparability with
 # older PyTorch.
 if "all_gather_into_tensor" not in dir(torch.distributed):
-    assert torch.distributed.is_available(), "PyTorch Distributed is Not available or Disabled."
+    if not torch.distributed.is_available():
+        raise RuntimeError("PyTorch Distributed is Not available or Disabled.")
     torch.distributed.all_gather_into_tensor = torch.distributed._all_gather_base
 
 def ensure_divisibility(numerator, denominator):