map location cpu

app.py

@@ -11,7 +11,7 @@ device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 #intialize the model
 model = models.model(pretrained=False, requires_grad=False).to(device)
 # load the model checkpoint
-checkpoint = torch.load('model.pth')
+checkpoint = torch.load('model.pth', map_location=device)
 # load model weights state_dict
 model.load_state_dict(checkpoint['model_state_dict'])
 model.eval()