updated model load command

app.py

@@ -7,10 +7,8 @@ from pytorch_grad_cam import GradCAM
 from pytorch_grad_cam.utils.image import show_cam_on_image
 from custom_resnet import Net
 
-# model = Net('batch')
-model = torch.load("model.pth", map_location=torch.device('cpu'))
-model.eval()
-model.to(torch.device('cpu'))
+model = Net('batch')
+model.load_state_dict(torch.load("model.pth", map_location=torch.device('cpu')), strict=False)
 
 classes = ('plane', 'car', 'bird', 'cat', 'deer',
            'dog', 'frog', 'horse', 'ship', 'truck')