Spaces:

nehalelkaref
/

flat-arabic-entity-classification

Build error

nehalelkaref commited on Nov 19, 2023

Commit

ed540bf

1 Parent(s): ac10e5c

Update network.py

Files changed (1) hide show

network.py CHANGED Viewed

@@ -55,7 +55,7 @@ class SpanNet(nn.Module):
     def load_model(cls, model_path, device=DEFAULT_DEVICE):
         res = torch.load(model_path, device)
         model = cls(**res['model_config'])
-        model.load_state_dict(res['model_state_dict'])
         model.eval()
         return model

     def load_model(cls, model_path, device=DEFAULT_DEVICE):
         res = torch.load(model_path, device)
         model = cls(**res['model_config'])
+        model.load_state_dict(res['model_state_dict'], strict=False)
         model.eval()
         return model