update predictor

predictor.py

@@ -198,11 +198,11 @@ def evaluation(model, epoch, device, dataloaders):
             print(preds)
 
 def inference(inp_img, classes = ['big', 'small'], epoch = 6):
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    device = torch.device("cpu")
     translator= Translator(to_lang="zh-TW")
 
     model = CUPredictor()
-    model.load_state_dict(torch.load(f'models/model_{epoch}.pt'))
+    model.load_state_dict(torch.load(f'models/model_{epoch}.pt', map_location=torch.device('cpu')))
     # load image-to-text model
     processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
     model_blip = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base")