Commit test

model.py

@@ -48,7 +48,7 @@ def main():
     # Load the final trained model (entire model, not just state_dict)
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     model_path = os.path.join(base_dir, "best_model.pth")
-    model = torch.load(model_path)
+    model = torch.load(model_path, map_location=device)
     # Also load the threshold from the best hyperparameters
     hyperparams_path = os.path.join(base_dir, "best_hyperparameters.json")
     with open(hyperparams_path, "r") as f: