fix test

neuralmagic · Oct 27, 2023 · ac68408 · ac68408
1 parent a56cf7f
commit ac68408
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/src/sparseml/core/lifecycle/session.py b/src/sparseml/core/lifecycle/session.py
@@ -42,7 +42,7 @@ class SparsificationLifecycle:
 
     def reset(self):
         for mod in self.modifiers:
-            if not mod.initialized_ or mod.finalized:
+            if not mod.initialized or mod.finalized:
                 continue
 
             try:

diff --git a/test_trainer.py b/test_trainer.py
@@ -1,6 +1,6 @@
-from sparseml.transformers.finetune.text_generation import main
-
 def run():
+    from sparseml.transformers.finetune.text_generation import main
+
     model = "./obcq_deployment"
     dataset_name = "wikitext"
     dataset_config_name = "wikitext-2-raw-v1"