Update app.py

app.py

@@ -2,10 +2,17 @@ import gradio as gr
 import torchvision.transforms as transforms
 from PIL import Image
 import torch
+from timm.models import create_model
 def predict(input_img):
     input_img =  Image.fromarray(np.uint8(input_img))
-    model1 = models.__dict__['resnet50'](num_classes=1)
-    model2 = models.__dict__['resnet50'](num_classes=1)
+    model1 = create_model(
+        'resnet50',
+        drop_rate=0.5,
+        num_classes=1,)
+    model2 = create_model(
+        'resnet50',
+        drop_rate=0.5,
+        num_classes=1,)
     
     loc = 'cuda:{}'.format(0)
     checkpoint1 = torch.load("./machine_full_best.tar", map_location=loc)