Merge branch 'main' into maanug/perf-hf-env-var

Author: maanug-nv

nemo/core/connectors/save_restore_connector.py

@@ -756,8 +756,23 @@ def _save_state_dict_to_disk(state_dict, filepath):
         torch.save(state_dict, filepath)
 
     @staticmethod
-    def _load_state_dict_from_disk(model_weights, map_location=None):
-        return torch.load(model_weights, map_location='cpu', weights_only=False)
+    def _load_state_dict_from_disk(model_weights, map_location='cpu'):
+        """
+        Load model state dict from disk.
+
+        Args:
+            model_weights: Path to the checkpoint file
+            map_location: Device to map tensors to
+
+        Returns:
+            State dict loaded from checkpoint
+
+        """
+        try:
+            return torch.load(model_weights, map_location=map_location, weights_only=True)
+        except Exception as e:
+            logging.error(f"Failed to load checkpoint with weights_only=True: {e}")
+            raise e
 
     @property
     def model_config_yaml(self) -> str:

requirements/requirements_lightning.txt

@@ -8,6 +8,6 @@ torchmetrics>=0.11.0
 transformers~=4.53.0
 wandb
 webdataset>=0.2.86
-nv_one_logger_core>=2.3.0
-nv_one_logger_training_telemetry>=2.3.0
-nv_one_logger_pytorch_lightning_integration>=2.3.0
+nv_one_logger_core>=2.3.1
+nv_one_logger_training_telemetry>=2.3.1
+nv_one_logger_pytorch_lightning_integration>=2.3.1