Update app.py

app.py

@@ -67,8 +67,9 @@ def server(input: Inputs, output: Outputs, session: Session):
     
     # Create an instance of the model architecture with the loaded configuration
     model = SamModel(config=model_config)
-    # Update the model by loading the weights from saved file.
-    model.load_state_dict(torch.load(str(dir / "checkpoint.pth")))
+    # Update the model by loading the weights from saved file
+    model_state_dict = torch.load(str(dir / "checkpoint.pth"), map_location=torch.device('cpu'))
+    model.load_state_dict(model_state_dict)
     
     # set the device to cuda if available, otherwise use cpu
     device = "cuda" if torch.cuda.is_available() else "cpu"