From d0abbf3bfd96ba17a2a1eb214f4dc222329298fb Mon Sep 17 00:00:00 2001
From: Rahul Tuli <rahul@neuralmagic.com>
Date: Mon, 21 Aug 2023 16:54:06 -0400
Subject: [PATCH] [Cherry Pick 1.5.4] Fix Index error with nm-transformers
 1.5.1 upgrade (#1708)

* Fix Index error with 1.5.1 transformers upgrade

* Deafult resolve_archive_file to None instead of empty list

* Update Hotfix version

* Remove failing tensorflow_v1 tests; they are outdated and not supported now
---
 .github/workflows/test-check.yaml             | 24 -------------------
 .../transformers/sparsification/trainer.py    |  4 ++--
 src/sparseml/version.py                       |  2 +-
 3 files changed, 3 insertions(+), 27 deletions(-)

diff --git a/.github/workflows/test-check.yaml b/.github/workflows/test-check.yaml
index c0488beda7d..b2f9c5ed37d 100644
--- a/.github/workflows/test-check.yaml
+++ b/.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:
diff --git a/src/sparseml/transformers/sparsification/trainer.py b/src/sparseml/transformers/sparsification/trainer.py
index 9ea0ce8bfc0..148be716253 100644
--- a/src/sparseml/transformers/sparsification/trainer.py
+++ b/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,
         )
diff --git a/src/sparseml/version.py b/src/sparseml/version.py
index 54680d62382..1e505895022 100644
--- a/src/sparseml/version.py
+++ b/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