Update app.py

app.py

@@ -41,7 +41,8 @@ tokenizer = AutoTokenizer.from_pretrained(MODEL_REPO)
 
 # Load model with appropriate dtype and device
 model_kwargs = {"torch_dtype": torch.float16} if device == "cuda" else {}
-model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_REPO, **model_kwargs).to(device)
+model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_REPO, **model_kwargs)
+model = model.to(device)  # Move model to device after full loading
 print(f"Model loaded on: {model.device}")
 
 # Ensure a pad token to avoid generate() quirks