Update app.py

app.py

@@ -16,7 +16,7 @@ transformer = transforms.Compose([
 ])
 
 
-model_name = 'custom'
+model_name = 'resnet'
 ### 2. Model and transforms preparation ###
 
 if model_name == 'custom':
@@ -36,7 +36,7 @@ elif model_name == 'resnet':
     # Load saved weights
     model.load_state_dict(
         torch.load(
-            f="./cnn-custom-model-version-4.pt",
+            f="./cnn-resnet-version-1.pt",
             map_location=torch.device("cpu"),  # load to CPU
         )
     )