Update app.py

app.py

@@ -5,7 +5,7 @@ import timm
 import gradio as gr
 import numpy as np
 
-device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+device = torch.device('cpu')
 
 labels = {0: 'bacterial_leaf_blight',
  1: 'bacterial_leaf_streak',
@@ -47,7 +47,7 @@ def predict(image = None) :
     image = np.transpose(image, (2, 0, 1))
     image = torch.tensor(image, dtype=torch.float32)
     model = timm.create_model('efficientnet_b0', pretrained=False, num_classes=10)
-    model.load_state_dict(torch.load("paddy_model.pth"))
+    model.load_state_dict(torch.load("paddy_model.pth", map_location=torch.device(device)))
     model.to(device)
 
     predicted = inference_fn(model, image)