Update app.py

app.py

@@ -6,6 +6,8 @@ import torchvision.transforms as transforms
 model = torchvision.models.resnet50(pretrained=True)
 model.fc = nn.Linear(model.fc.in_features, 2)
 model.load_state_dict(torch.load("model.pth"))
+input_batch = input_batch.to('cpu')
+model.to('cpu')
 model.eval()
 
 transform = transforms.Compose([
@@ -30,10 +32,7 @@ def inference(input_image):
     input_tensor = preprocess(input_image)
     input_batch = input_tensor.unsqueeze(0) # create a mini-batch as expected by the model
 
-    # move the input and model to GPU for speed if available
-    if torch.cuda.is_available():
-        input_batch = input_batch.to('cuda')
-        model.to('cuda')
+        
 
     with torch.no_grad():
         output = model(input_batch)