Put train.py

train.py

@@ -148,6 +148,8 @@ def main():
         push_to_hub=False,
         hub_model_id=args.output,
     )
+    if args.max_steps > 0:
+        cfg.max_steps = args.max_steps
 
     trainer = SFTTrainer(
         model=model,
@@ -157,10 +159,7 @@ def main():
         processing_class=tokenizer,
         peft_config=peft,
     )
-    train_kwargs = {}
-    if args.max_steps > 0:
-        train_kwargs["max_steps"] = args.max_steps
-    trainer.train(**train_kwargs)
+    trainer.train()
 
     if args.smoke:
         print("SMOKE OK")