update predictor

predictor.py

@@ -201,7 +201,7 @@ def inference(inp_img, classes = ['big', 'small'], epoch = 6):
     device = torch.device("cuda")
     translator= Translator(to_lang="zh-TW")
 
-    model = CUPredictor().to(device)
+    model = CUPredictor()
     model.load_state_dict(torch.load(f'models/model_{epoch}.pt'))
     # load image-to-text model
     processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
@@ -218,13 +218,13 @@ 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.to(device)
+        inputs = image_tensor
         outputs_c, outputs_h, outputs_b, outputs_w, outputs_hi = model(inputs)
         _, preds = torch.max(outputs_c, 1)
         idx = preds.numpy()[0]
 
         # unconditional image captioning
-        inputs = processor(inp_img, return_tensors="pt").to(device)
+        inputs = processor(inp_img, return_tensors="pt")
         out = model_blip.generate(**inputs)
         description = processor.decode(out[0], skip_special_tokens=True)
         description_tw = translator.translate(description)