Update app.py

app.py

@@ -14,7 +14,7 @@ with open("imagenet_classes.json", "r") as f:
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model = resnet50(num_classes=1000, drop_path_rate=0.0, use_blurpool=True)
 # model.load_state_dict(torch.load("best_resnet50_imagenet_1k.pt", map_location=device))
-checkpoint = torch.load('best_resnet50_imagenet_1k.pt')
+checkpoint = torch.load('best_resnet50_imagenet_1k.pt', map_location=device)
 model.load_state_dict(checkpoint['model_state_dict'])
 model.to(device)
 model.eval()