fix bugs

main.py

@@ -92,12 +92,13 @@ def predict_image(image_url):
     img_cp = np.float32(img_cp) / 255
     input_tensor = preprocess_image(img_cp, mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
     input_tensor = torch.Tensor(input_tensor)
-    input_tensor.cuda()
+    # input_tensor.cuda()
     
     output = model(input_tensor)
     # print(torch.max(output, 1))
     _, preds_tensor = torch.max(output, 1)
-    preds = np.squeeze(preds, preds_tensor.numpy()) if not use_cuda else np.squeeze(preds_tensor.cpu().numpy())
+    # preds = np.squeeze(preds, preds_tensor.numpy()) if not use_cuda else np.squeeze(preds_tensor.cpu().numpy())
+    preds = np.squeeze(preds_tensor.cpu().numpy()) 
     print(preds)
 
     class_name = classes[preds]
@@ -108,7 +109,7 @@ def predict_image(image_url):
         img = cv2.resize(img, (224, 224))
         img = np.float32(img) / 255
         input_tensor = torch.Tensor(input_tensor)
-        input_tensor.cuda()
+        # input_tensor.cuda()
         
         input_tensor = preprocess_image(img, mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
         targets = [ClassifierOutputTarget(0)]