Fixed map location

app.py

@@ -12,7 +12,7 @@ print('Creating the model')
 model = get_model('checkpoint.ckpt')
 print('Loading the classifier')
 classifier = CustomEfficientNet(target_size=7, pretrained=False)
-classifier.load_state_dict(torch.load('class_efficientB0_taco_7_class.pth'))
+classifier.load_state_dict(torch.load('class_efficientB0_taco_7_class.pth', map_location='cpu'))
 
 def plot_img_no_mask(image, boxes, labels):
     colors = {