Update app.py

app.py

@@ -24,8 +24,12 @@ from torchvision import transforms
 _DEFAULT_IMAGE_TENSOR_NORMALIZATION_MEAN = [0.485, 0.456, 0.406]
 _DEFAULT_IMAGE_TENSOR_NORMALIZATION_STD = [0.229, 0.224, 0.225]
 
-def evaluate(img, model, id2entity, target_list, args):
+def evaluate(img, id2entity, target_list, args):
+    num_ent_id = len(id2entity)
+    model = DistMult(args, num_ent_id, target_list, args.device)
+    model.to(args.device)
     model.eval()
+    
     torch.set_grad_enabled(False)
 
     overall_id_to_name = json.load(open('overall_id_to_name.json'))
@@ -146,10 +150,6 @@ if __name__=='__main__':
     
     target_list = generate_target_list(datacsv, entity2id)
 
-    model = DistMult(args, num_ent_id, target_list, args.device)
-
-    model.to(args.device)
-
     # restore from ckpt
     if args.ckpt_path:
         ckpt = torch.load(args.ckpt_path, map_location=args.device)
@@ -158,7 +158,7 @@ if __name__=='__main__':
 
     species_model = gr.Interface(
         evaluate,
-        [gr.inputs.Image(shape=(200, 200)), model, id2entity, target_list, args],
+        [gr.inputs.Image(shape=(200, 200)), id2entity, target_list, args],
         outputs="label",
         title = 'Species Classification',
     )