Update app.py

app.py

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