Update app.py

app.py

@@ -34,7 +34,7 @@ def load_the_model():
         x = self.fc3(x)
         return x
     
-  loaded_model = Net().to(device)
+  loaded_model = Net()#.to(device)
   loaded_model.load_state_dict(torch.load(path_of_model))
   loaded_model.eval()
   return loaded_model
@@ -86,7 +86,7 @@ def predict_on_base64(model, base64_file):
         custom_image_transformed_with_batch_size = custom_image_transformed.unsqueeze(dim=0)
         
         # Make a prediction on image with an extra dimension
-        custom_image_pred = model(custom_image_transformed_with_batch_size.to(device))
+        custom_image_pred = model(custom_image_transformed_with_batch_size) # .to(device)
         # Getting the probs
         prob = torch.softmax(custom_image_pred, dim=1)
         # Getting the sample prob