Revised saving

2025-07-23 15:17:14 +08:00
parent 8a7a011cb1
commit bace84f9e9
2 changed files with 7 additions and 2 deletions
--- a/utils/checkpoint.py
+++ b/utils/checkpoint.py
@@ -86,7 +86,7 @@ class CheckpointManager:
                model_path = os.path.join(temp_dir, 'pytorch_model.bin')
                # Add extra safety check for model state dict
                model_state = state_dict['model_state_dict']
-                if not model_state:
+                if model_state is None or (isinstance(model_state, dict) and len(model_state) == 0):
                    raise RuntimeError("Model state dict is empty")
                
                self._safe_torch_save(model_state, model_path)