fix inference test

flexflow · Nov 4, 2024 · fc884fe · fc884fe
1 parent d09ba0c
commit fc884fe
Showing 1 changed file with 9 additions and 9 deletions.
diff --git a/tests/peft/hf_utils.py b/tests/peft/hf_utils.py
@@ -223,15 +223,15 @@ def save_lora_weights(self, model, pre_finetuning=False):
         if not pre_finetuning:
             self.step_count += 1
 
-    def on_step_end(
-        self, args, state, control, model, tokenizer, optimizer, lr_scheduler, **kwargs
-    ):
-        self.save_lora_weights(model, pre_finetuning=False)
-
-    def on_step_begin(
-        self, args, state, control, model, tokenizer, optimizer, lr_scheduler, **kwargs
-    ):
-        self.save_lora_weights(model, pre_finetuning=True)
+    def on_step_end(self, args, state, control, **kwargs):
+        model_ = kwargs.get("model", None)
+        assert model_ is not None
+        self.save_lora_weights(model_, pre_finetuning=False)
+
+    def on_step_begin(self, args, state, control, **kwargs):
+        model_ = kwargs.get("model", None)
+        assert model_ is not None
+        self.save_lora_weights(model_, pre_finetuning=True)
 
     def on_train_end(self, args, state, control, **kwargs):
         if verbose: