Spaces:

saritha
/

sugarcanedisease

Sleeping

saritha commited on Jan 1

Commit

3a0ccc3

verified ·

1 Parent(s): 17d4c8d

Update app.py

Files changed (1) hide show

app.py CHANGED Viewed

@@ -26,7 +26,7 @@ def suppress_stdout():
 # Load the saved model and suppress the warnings
 with suppress_stdout():
     model = ViTForImageClassification.from_pretrained('google/vit-base-patch16-224-in21k', num_labels=6)
-    model.load_state_dict(torch.load('vit_sugarcane_disease_detection.pth'))
     model.eval()
 # Define the same transformation used during training

 # Load the saved model and suppress the warnings
 with suppress_stdout():
     model = ViTForImageClassification.from_pretrained('google/vit-base-patch16-224-in21k', num_labels=6)
+    model.load_state_dict(torch.load('vit_sugarcane_disease_detection.pth', map_location=torch.device('cpu')))
     model.eval()
 # Define the same transformation used during training