Skip processing if extra_state loads as None (#11778)

Signed-off-by: Jan Lasek <[email protected]>
NVIDIA · Jan 8, 2025 · ee12d88 · ee12d88
1 parent eb9848b
commit ee12d88
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/nemo/export/trt_llm/nemo_ckpt_loader/nemo_file.py b/nemo/export/trt_llm/nemo_ckpt_loader/nemo_file.py
@@ -93,6 +93,9 @@ def preprocess_scaling_factors_for_local_export(state_dict: Dict[str, Any]) -> D
     scales = {}
 
     for key, value in scales_dict.items():
+        if value is None:
+            continue
+
         value.seek(0)
         extra_state = torch.load(value)
         if extra_state is not None and 'scale_fwd' in extra_state: