updated model loading function to support zsbert model trained with o…

…ld transformer library (4.21.3) (#61) Signed-off-by: Marcos Martinez <[email protected]> Co-authored-by: Alberto Purpura <[email protected]>
IBM · Aug 24, 2023 · f5fcf7f · f5fcf7f
1 parent ebaf046
commit f5fcf7f
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/zshot/relation_extractor/zsrc/zero_shot_rel_class.py b/zshot/relation_extractor/zsrc/zero_shot_rel_class.py
@@ -17,7 +17,7 @@ def load_model(device: Optional[Union[str, torch.device]] = None):
     if not os.path.isfile(MODEL_PATH):
         download_file(MODEL_REMOTE_URL, MODELS_CACHE_PATH)
 
-    model.load_state_dict(torch.load(MODEL_PATH))
+    model.load_state_dict(torch.load(MODEL_PATH), strict=False)
     model.to(device)
     model.eval()
     return model