update predictor

predictor.py

@@ -198,7 +198,7 @@ def evaluation(model, epoch, device, dataloaders):
             print(preds)
 
 def inference(inp_img, classes = ['big', 'small'], epoch = 6):
-    device = torch.device("cuda")
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     translator= Translator(to_lang="zh-TW")
 
     model = CUPredictor()
@@ -218,7 +218,7 @@ def inference(inp_img, classes = ['big', 'small'], epoch = 6):
     image_tensor = trans(inp_img)
     image_tensor = image_tensor.unsqueeze(0)
     with torch.no_grad():
-        inputs = image_tensor
+        inputs = image_tensor.to(device)
         outputs_c, outputs_h, outputs_b, outputs_w, outputs_hi = model(inputs)
         _, preds = torch.max(outputs_c, 1)
         idx = preds.numpy()[0]