add flax mbart in auto seq2seq lm (#13560)

fc3551a6 · Suraj Patil · GitHub · 3081d386 · fc3551a6
Unverified Commit fc3551a6 authored Sep 14, 2021 by Suraj Patil Committed by GitHub Sep 14, 2021
Show whitespace changes
Inline Side-by-side

Showing with 1 addition and 0 deletions

src/transformers/models/auto/modeling_flax_auto.py src/transformers/models/auto/modeling_flax_auto.py +1 -0

No files found.
--- a/src/transformers/models/auto/modeling_flax_auto.py
+++ b/src/transformers/models/auto/modeling_flax_auto.py
@@ -81,6 +81,7 @@ FLAX_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING_NAMES = OrderedDict(
    [
        # Model for Seq2Seq Causal LM mapping
        ("bart", "FlaxBartForConditionalGeneration"),
+        ("mbart", "FlaxMBartForConditionalGeneration"),
        ("t5", "FlaxT5ForConditionalGeneration"),
        ("mt5", "FlaxMT5ForConditionalGeneration"),
        ("marian", "FlaxMarianMTModel"),