Update app.py

app.py

@@ -10,7 +10,7 @@ from huggingface_hub import hf_hub_download
 
 tokenizer = AutoTokenizer.from_pretrained('bert-base-uncased', )
 
-ckpt = torch.load(hf_hub_download('tennant/MUG', 'mae_bert_vit_b_cc3m.pth'), map_location='cpu')
+ckpt = torch.load(hf_hub_download('tennant/MUG', 'laion_mug_vit_base_5ep.pth'), map_location='cpu')
 
 new_dict = OrderedDict()
 for k, v in ckpt.items():
@@ -19,7 +19,8 @@ for k, v in ckpt.items():
 
 model = mae_vit_base_patch16(uni_dim=768, less_u=True)
 
-model.load_state_dict(new_dict)
+msg = model.load_state_dict(new_dict, strict=False)
+print(msg)
 if torch.cuda.is_available():
     model.cuda()
 model.eval()