Add comments

Signed-off-by: Guyue Huang <[email protected]>
NVIDIA · Feb 22, 2025 · f6ccf80 · f6ccf80
1 parent fd76d53
commit f6ccf80
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 0 deletions.
diff --git a/transformer_engine/pytorch/tensor/float8_tensor.py b/transformer_engine/pytorch/tensor/float8_tensor.py
@@ -481,7 +481,10 @@ def _set_data(self, tensor: torch.Tensor) -> None:
         casts to FP8.
 
         """
+
+        # Synchronize to avoid a race condition
         torch.cuda.synchronize()
+
         # Tensor device
         new_device = tensor.device if tensor.is_cuda else self.device
 

diff --git a/transformer_engine/pytorch/tensor/mxfp8_tensor.py b/transformer_engine/pytorch/tensor/mxfp8_tensor.py
@@ -366,6 +366,9 @@ def _set_data(self, tensor: torch.Tensor) -> None:
 
         """
 
+        # Synchronize to avoid a race condition
+        torch.cuda.synchronize()
+
         # Tensor device
         new_device = tensor.device if tensor.is_cuda else self.device