Update app.py

app.py

@@ -5,11 +5,12 @@ import gradio as gr
 from PIL import Image
 from pytorch_grad_cam import GradCAM
 from pytorch_grad_cam.utils.image import show_cam_on_image
-from resnet import ResNet18
+from resnet import custom_ResNet
 import gradio as gr
 
-model = ResNet18()
+model = custom_ResNet()
 model.load_state_dict(torch.load("custom_resnet_model.pth", map_location=torch.device('cpu')), strict=False)
+model.setup(stage="test")
 
 inv_normalize = transforms.Normalize(
     mean=[-0.50/0.23, -0.50/0.23, -0.50/0.23],