Update app.py

app.py

@@ -6,10 +6,9 @@ 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_lightning import ResNet18Model
 import gradio as gr
 
-model = ResNet18Model.load_from_checkpoint("epoch=19-step=3920.ckpt")
+model = torch.load('best.pt')
 
 inv_normalize = transforms.Normalize(
     mean = [-0.50/0.23, -0.50/0.23, -0.50/0.23],