Spaces:

PedroMartelleto
/

breast-cancer-ultrasound-explained

Runtime error

PedroMartelleto commited on Dec 20, 2022

Commit

87fbe80

•

1 Parent(s): 56146b3

small fixes

Files changed (1) hide show

app.py CHANGED Viewed

@@ -68,7 +68,7 @@ def create_model_from_checkpoint():
     # Loads a model from a checkpoint
     model = resnet50()
     model.fc = nn.Linear(model.fc.in_features, 3)
-    model.load_state_dict(torch.load("best_model"))
     model.eval()
     return model

     # Loads a model from a checkpoint
     model = resnet50()
     model.fc = nn.Linear(model.fc.in_features, 3)
+    model.load_state_dict(torch.load("best_model", map_location=torch.device('cpu')))
     model.eval()
     return model