add docstring

Signed-off-by: Xin Yao <[email protected]>
NVIDIA · Oct 21, 2024 · 0987833 · 0987833
1 parent 3fe0c0f
commit 0987833
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/transformer_engine/pytorch/utils.py b/transformer_engine/pytorch/utils.py
@@ -312,9 +312,11 @@ def devices_match(device1: torch.device, device2: torch.device) -> bool:
 if TORCH_MAJOR == 2 and TORCH_MINOR >= 4:
 
     def torch_get_autocast_gpu_dtype() -> torch.dtype:
+        """Get PyTorch autocast GPU dtype."""
         return torch.get_autocast_dtype("cuda")
 
 else:
 
     def torch_get_autocast_gpu_dtype() -> torch.dtype:
+        """Get PyTorch autocast GPU dtype."""
         return torch.get_autocast_gpu_dtype()