Update app.py

app.py

@@ -13,7 +13,9 @@ with open("imagenet_classes.json", "r") as f:
 # Load model
 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))
+# model.load_state_dict(torch.load("best_resnet50_imagenet_1k.pt", map_location=device))
+checkpoint = torch.load('best_resnet50_imagenet_1k.pt')
+model.load_state_dict(checkpoint['model_state_dict'])
 model.to(device)
 model.eval()