Spaces:

Vrk
/

SeeFood

Sleeping

Vrk commited on Dec 14, 2022

Commit

edb6bfb

1 Parent(s): 5fcdc39

Update utils.py

Files changed (1) hide show

utils.py CHANGED Viewed

@@ -76,7 +76,7 @@ def model_pred_tf(model_path, img, class_names=classes):
 def get_model_pt(model_path):
     model = timm.create_model('vit_base_patch16_224', pretrained=False)
     model.head = nn.Linear(in_features=768, out_features=len(classes), bias=True)
-    model.load_state_dict(torch.load('models/ViT-101-1.pt', map_location='cpu'))
     return model
 def load_prepare_image_pt(input_image):

 def get_model_pt(model_path):
     model = timm.create_model('vit_base_patch16_224', pretrained=False)
     model.head = nn.Linear(in_features=768, out_features=len(classes), bias=True)
+    model.load_state_dict(torch.load(model_path, map_location='cpu'))
     return model
 def load_prepare_image_pt(input_image):