[Cherry Pick 1.5.4] Fix Index error with nm-transformers 1.5.1 upgrade

.github/workflows/test-check.yaml

@@ -160,30 +160,6 @@ jobs:
         run: pip3 install .[dev,torchvision,onnxruntime]
       - name: "🔬 Running pytorch tests"
         run: make test TARGETS=pytorch
-  tensorflow-v1-tests:
-    runs-on: ubuntu-22.04
-    env:
-      SPARSEZOO_TEST_MODE: "true"
-    needs: test-setup
-    if:  ${{needs.test-setup.outputs.tensorflow_v1 == 1}}
-    steps:
-      - uses: actions/setup-python@v4
-        with:
-          python-version: '3.7'
-      - uses: actions/checkout@v2
-      - uses: actions/checkout@v2
-        with:
-          repository: "neuralmagic/sparsezoo"
-          path: "sparsezoo"
-          ref: ${{needs.test-setup.outputs.branch}}
-      - name: "⚙️ Install sparsezoo dependencies"
-        run: pip3 install -U pip && pip3 install setuptools sparsezoo/
-      - name: "Clean sparsezoo directory"
-        run: rm -r sparsezoo/
-      - name: "⚙️ Install dependencies"
-        run: pip3 install .[dev,tf_v1,onnxruntime]
-      - name: "🔬 Running tensorflow_v1 tests"
-        run: make test TARGETS=tensorflow_v1
   compat-pytorch-1_9-pytorch-tests:
     runs-on: ubuntu-22.04
     env:

src/sparseml/transformers/sparsification/trainer.py

@@ -673,11 +673,11 @@ def _reload_model_state(self, load_path: str, orig_state_dict: Dict[str, Any]):
         loaded_state_dict = torch.load(
             os.path.join(load_path, WEIGHTS_NAME), map_location="cpu"
         )
-        _, missing, unexpected, _, _ = self.model._load_pretrained_model(
+        _, missing, unexpected, _, _, _ = self.model._load_pretrained_model(
             model=self.model,
             state_dict=loaded_state_dict,
             loaded_keys=list(loaded_state_dict.keys()),
-            resolved_archive_file=[],
+            resolved_archive_file=None,
             pretrained_model_name_or_path=load_path,
             _fast_init=False,
         )

src/sparseml/version.py

@@ -19,7 +19,7 @@
 from datetime import date
 
 
-version_base = "1.5.3"
+version_base = "1.5.4"
 is_release = False  # change to True to set the generated version as a release version