Update app.py

app.py

@@ -191,27 +191,29 @@ def find_example_files():
 # ==========================
 
 BASE_MODEL_ID = "Salesforce/codet5p-770m"
-FINETUNED_REPO_ID = "OSS-forge/codet5p-770m-pyresbugs"   
+FINETUNED_REPO_ID = "OSS-forge/codet5p-770m-pyresbugs"  
 FINETUNED_FILENAME = "pytorch_model.bin"
 
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
 print(f"Loading tokenizer from base model: {BASE_MODEL_ID}")
 tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL_ID)
 
 print(f"Loading base model: {BASE_MODEL_ID}")
 model = AutoModelForSeq2SeqLM.from_pretrained(BASE_MODEL_ID)
+model.to(device)
 
 print(f"Downloading fine-tuned weights from repo: {FINETUNED_REPO_ID}")
 ckpt_path = hf_hub_download(FINETUNED_REPO_ID, FINETUNED_FILENAME)
 
 print(f"Loading state_dict from: {ckpt_path}")
-state_dict = torch.load(ckpt_path, map_location=torch.device("cpu"))
+state_dict = torch.load(ckpt_path, map_location="cpu")
 
 if "model_state_dict" in state_dict:
     state_dict = state_dict["model_state_dict"]
 
 missing, unexpected = model.load_state_dict(state_dict, strict=False)
-print("Loaded fine-tuned weights.")
-print("Missing keys:", len(missing), "Unexpected keys:", len(unexpected))
+print(f"Loaded fine-tuned weights. Missing keys: {len(missing)}, unexpected keys: {len(unexpected)}")
 
 model.eval()