yxchng commited on
Commit
4f646a4
1 Parent(s): 46907a9
Files changed (1) hide show
  1. app.py +1 -1
app.py CHANGED
@@ -197,7 +197,7 @@ model = WrapperModel(single_model.backbone, single_bert_model, maskformer_head)
197
 
198
  checkpoint = torch.load(weights, map_location='cpu')
199
 
200
- model.load_state_dict(checkpoint, strict=False)
201
  model.to(device)
202
  model.eval()
203
  #single_bert_model.load_state_dict(checkpoint['bert_model'])
 
197
 
198
  checkpoint = torch.load(weights, map_location='cpu')
199
 
200
+ model.load_state_dict(checkpoint['model'], strict=False)
201
  model.to(device)
202
  model.eval()
203
  #single_bert_model.load_state_dict(checkpoint['bert_model'])