Spaces:

kakasher
/

Cat-Safe-Plant-Classfier

Running

kakasher commited on Sep 25, 2024

Commit

8531995

1 Parent(s): a00a922

changed device to cpu

Files changed (1) hide show

app.py CHANGED Viewed

@@ -36,7 +36,7 @@ model = None
 # Load model
 try:
     model = timm.create_model('vit_base_patch16_224', pretrained=False, num_classes=47)
-    model.load_state_dict(torch.load(model_path, weights_only=False))
     model.eval()
 except Exception as e:
     print(f"Error loading model: {str(e)}")

 # Load model
 try:
     model = timm.create_model('vit_base_patch16_224', pretrained=False, num_classes=47)
+    model.load_state_dict(torch.load(model_path, weights_only=False, map_location=torch.device('cpu')))
     model.eval()
 except Exception as e:
     print(f"Error loading model: {str(e)}")