make size of tensor even integer.

argonne-lcf · Oct 5, 2024 · 9f12a50 · 9f12a50
1 parent b2bf926
commit 9f12a50
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/dlio_benchmark/checkpointing/base_checkpointing.py b/dlio_benchmark/checkpointing/base_checkpointing.py
@@ -54,7 +54,7 @@ def __init__(self, ext):
                 self.layer_state = dict()
                 for index, state in enumerate(self.args.layer_parameters):
                     if state > 0:
-                        self.layer_state[str(index)] = self.get_tensor(state / self.args.tensor_parallelism)
+                        self.layer_state[str(index)] = self.get_tensor(state // self.args.tensor_parallelism)
 
     @abstractmethod
     def get_tensor(self, size):