Update app.py

app.py

@@ -6,6 +6,8 @@ import torch.nn as nn
 import pandas as pd
 import numpy as np 
 import gradio as gr
+import warnings
+warnings.filterwarnings("ignore", category=UserWarning) 
 
 class_names = ['apple_pie',
  'bibimbap',
@@ -41,7 +43,8 @@ def predict(img_path):
     # Apply transformations
     image = test_transforms(image)
 
-    inf_model = models.resnet18(pretrained=False)
+    inf_model = models.resnet18(pretrained=False, weights='ResNet50_Weights.DEFAULT')
+
     num_ftrs = inf_model.fc.in_features
     # Here the size of each output sample is set to 2.
     # Alternatively, it can be generalized to nn.Linear(num_ftrs, len(class_names)).