fixed datatype issue

argonne-lcf · Feb 7, 2025 · 9690df2 · 9690df2
1 parent 843cabb
commit 9690df2
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 8 deletions.
diff --git a/dlio_benchmark/checkpointing/base_checkpointing.py b/dlio_benchmark/checkpointing/base_checkpointing.py
@@ -140,9 +140,9 @@ def __init__(self, ext):
             logging.info(f"{utcnow()} Layer size: {ss} GB")
             logging.info(f"{utcnow()} Optimizer state size: {opt} GB")
             logging.info(f"{utcnow()} Total checkpoint size: {self.checkpoint_size} GB")
-            
+
     @abstractmethod
-    def get_tensor(self, size):
+    def get_tensor(self, size, datatype="int8"):
         return []
 
     @abstractmethod

diff --git a/dlio_benchmark/checkpointing/pytorch_checkpointing.py b/dlio_benchmark/checkpointing/pytorch_checkpointing.py
@@ -25,18 +25,21 @@
 from dlio_benchmark.utils.utility import DLIOMPI
 
 def get_torch_datatype(datatype):
-    if datatype == "f32":
+    if datatype == "fp32":
         return torch.float32
-    if datatype == "f16":
+    elif datatype == "fp16":
         return torch.float16
-    if datatype == "f64":
+    elif datatype == "fp64":
         return torch.float64
-    if datatype == "int8":
+    elif datatype == "int8":
         return torch.int8
-    if datatype == "uint8":
+    elif datatype == "uint8":
         return torch.uint8
-    if datatype == "bf16": # bfloat16
+    elif datatype == "bf16": # bfloat16
         return torch.bfloat16
+    else:
+        raise Exception(f"Invalid datatype {datatype}")
+
 
 dlp = Profile(MODULE_CHECKPOINT)