Update Predict.py

Predict.py

@@ -5,16 +5,14 @@ from PIL import Image
 from torchvision import transforms as T
 import torch 
 
-
-checkpoint = torch.load('last_checkpoint.pt')
+device = torch.device('cpu')
+checkpoint = torch.load('last_checkpoint.pt', map_location = device)
 with open('answer.json', 'r', encoding = 'utf8') as f:
     answer_space = json.load(f)
 swap_space = {v : k for k, v in answer_space.items()}
 
-
-device = torch.device('cpu')
 model = BaseModel().to(device)
-model.load_state_dict(checkpoint['model_state_dict'], map_locaption = device)
+model.load_state_dict(checkpoint['model_state_dict'])
 
 def generate_caption(image, question):
     if isinstance(image, np.ndarray):