cp: check asr models (14989) into r2.5.0 (#15017)

* check asr models (#14989)

* check asr models

Signed-off-by: nithinraok <nithinrao.koluguri@gmail.com>

* Apply isort and black reformatting

Signed-off-by: nithinraok <nithinraok@users.noreply.github.com>

* update return

Signed-off-by: nithinraok <nithinrao.koluguri@gmail.com>

---------

Signed-off-by: nithinraok <nithinrao.koluguri@gmail.com>
Signed-off-by: nithinraok <nithinraok@users.noreply.github.com>
Co-authored-by: nithinraok <nithinraok@users.noreply.github.com>
Co-authored-by: Charlie Truong <chtruong@nvidia.com>
Signed-off-by: NeMo Bot <nemo-bot@nvidia.com>

* Pass timeout when running speech functional tests (#15012)

Signed-off-by: Charlie Truong <chtruong@nvidia.com>

* Add timeout to speech test

Signed-off-by: Charlie Truong <chtruong@nvidia.com>

* Use cpu runner for asr cpu unit tests

Signed-off-by: Charlie Truong <chtruong@nvidia.com>

---------

Signed-off-by: nithinraok <nithinrao.koluguri@gmail.com>
Signed-off-by: nithinraok <nithinraok@users.noreply.github.com>
Signed-off-by: NeMo Bot <nemo-bot@nvidia.com>
Signed-off-by: Charlie Truong <chtruong@nvidia.com>
Co-authored-by: Nithin Rao <nithinrao.koluguri@gmail.com>
Co-authored-by: nithinraok <nithinraok@users.noreply.github.com>

Author: chtruong814

.github/workflows/cicd-main-speech.yml

@@ -39,7 +39,7 @@ jobs:
             runner: self-hosted-azure-gpus-1
             timeout: 30
           - script: L0_Unit_Tests_CPU_ASR
-            runner: self-hosted-azure-cpu
+            runner: azure-gpu-vm-runner1-cpu
             cpu-only: true
             timeout: 20
           - script: L0_Unit_Tests_GPU_TTS
@@ -180,6 +180,7 @@ jobs:
             script: SPEECHLM_HF_Training_DuplexS2SSpeechDecoder
           - runner: self-hosted-azure
             script: SPEECHLM_HF_Training_SALM
+            timeout: 20
     needs: [unit-tests]
     runs-on: ${{ matrix.runner }}
     name: ${{ matrix.is-optional && 'PLEASEFIXME_' || '' }}${{ matrix.script }}
@@ -195,4 +196,5 @@ jobs:
           script: ${{ matrix.script }}
           tests_to_run: ${{ inputs.test_to_run }}
           image: ${{ inputs.image-name }}
+          timeout: ${{ matrix.timeout || 10 }}
           is_optional: ${{ matrix.is-optional || false }}

nemo/core/connectors/save_restore_connector.py

@@ -754,8 +754,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: