Merge pull request #5668 from siddhu001/Fix_transformers_load

Fix loading pre-trained model from transformers
espnet · Feb 21, 2024 · 6834444 · 6834444
2 parents a50d6a0 + 8ca64cb
commit 6834444
Showing 1 changed file with 14 additions and 2 deletions.
diff --git a/espnet2/tasks/abs_task.py b/espnet2/tasks/abs_task.py
@@ -2089,8 +2089,20 @@ def build_model_from_file(
                         }
                         model.load_state_dict(state_dict, strict=not use_lora)
                     else:
-                        raise
+                        if any(["postdecoder" in k for k in state_dict.keys()]):
+                            model.load_state_dict(
+                                state_dict,
+                                strict=False,
+                            )
+                        else:
+                            raise
                 else:
-                    raise
+                    if any(["postdecoder" in k for k in state_dict.keys()]):
+                        model.load_state_dict(
+                            state_dict,
+                            strict=False,
+                        )
+                    else:
+                        raise
 
         return model, args