added sanity check

argonne-lcf · Feb 21, 2025 · d695c60 · d695c60
1 parent aafa316
commit d695c60
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 0 deletions.
diff --git a/dlio_benchmark/checkpointing/pytorch_checkpointing.py b/dlio_benchmark/checkpointing/pytorch_checkpointing.py
@@ -73,7 +73,10 @@ def save_state(self, suffix, state, fsync = False):
     @dlp.log
     def load_state(self, suffix, state):
         name = self.get_name(suffix)
+        state = dict() # clear up
         state = torch.load(name)
+        logging.debug(f"checkpoint state loaded: {state}")
+        assert(len(state.keys())>0)
 
     @dlp.log
     def save_checkpoint(self, epoch, step_number):

diff --git a/dlio_benchmark/checkpointing/tf_checkpointing.py b/dlio_benchmark/checkpointing/tf_checkpointing.py
@@ -71,7 +71,10 @@ def save_state(self, suffix, state, fsync = False):
     @dlp.log
     def load_state(self, suffix, state):
         name = self.get_name(suffix)
+        state = dict() # clear up
         state = tf.train.load_checkpoint(name)
+        logging.debug(f"checkpoint state loaded: {state}")
+        assert(len(state.keys)!=0)
 
     @dlp.log
     def save_checkpoint(self, epoch, step_number):