Update app.py

app.py

@@ -69,8 +69,8 @@ def load_checkpoint_for_inference(filepath, model_class):
     # (You need the Model Class definition handy)
     model = model_class()
 
-    # Load the checkpoint file
-    checkpoint = torch.load(filepath)
+    # Load the checkpoint file with map_location to handle device mismatch
+    checkpoint = torch.load(filepath, map_location=torch.device('cpu'))
 
     # Load the state dictionary into the model instance
     model.load_state_dict(checkpoint)
@@ -78,8 +78,7 @@ def load_checkpoint_for_inference(filepath, model_class):
     # Set the model to evaluation mode for inference
     model.eval()
 
-    # Optional: Move the model to the appropriate device (GPU/CPU)
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    # Move the model to the device (GPU/CPU)
     model.to(device)
 
     print(f"Checkpoint loaded.")