correctly handle mt5 (huggingface#9879)

Qbiwan · Jan 31, 2021 · daddaf7 · daddaf7
1 parent 8818b6c
commit daddaf7
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/examples/seq2seq/utils.py b/examples/seq2seq/utils.py
@@ -563,7 +563,7 @@ def freeze_embeds(model):
     """Freeze token embeddings and positional embeddings for bart, just token embeddings for t5."""
     model_type = model.config.model_type
 
-    if model_type == "t5":
+    if model_type in ["t5", "mt5"]:
         freeze_params(model.shared)
         for d in [model.encoder, model.decoder]:
             freeze_params(d.embed_tokens)