Update fine_tuning.py

fine_tuning.py

@@ -24,6 +24,7 @@ os.environ["WANDB_MODE"] = "disabled"
 # -----------------------------
 device = "cuda" if torch.cuda.is_available() else "cpu"
 tokenizer = transformers.AutoTokenizer.from_pretrained("TinyLlama/TinyLlama-1.1B-Chat-v1.0")
+model = None
 
 # -----------------------------
 # LoRA / MoE Modules
@@ -157,7 +158,6 @@ def load_and_train(model_id="TinyLlama/TinyLlama-1.1B-Chat-v1.0"):
     print("Training started")
     trainer.train()
     model.eval()
-    return model, tokenizer, device
 
 # -----------------------------
 # GENERATE ANSWER