allow bf16 transfer dtype in bench script

PrimeIntellect-ai · Oct 4, 2024 · b290a76 · b290a76
1 parent 5d0e704
commit b290a76
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/scripts/all_reduce_test.py b/scripts/all_reduce_test.py
@@ -13,13 +13,15 @@
 class TorchDtype(str, Enum):
     FLOAT32 = "float32"
     FLOAT16 = "float16"
+    BFLOAT16 = "bfloat16"
     UINT8 = "uint8"
 
 
 TORCH_DTYPE_MAP = {
     None: None,
     TorchDtype.FLOAT32: torch.float32,
     TorchDtype.FLOAT16: torch.float16,
+    TorchDtype.BFLOAT16: torch.bfloat16,
     TorchDtype.UINT8: torch.uint8,
 }