FlaxAutoModelForSeq2SeqLM bug fixed

run_translation_t5_flax.py

@@ -50,6 +50,7 @@ from transformers import (
     FlaxAutoModelForSeq2SeqLM,
     HfArgumentParser,
     TrainingArguments,
+    FlaxT5ForConditionalGeneration,
     is_tensorboard_available,
 )
 from transformers.file_utils import is_offline_mode
@@ -374,11 +375,11 @@ def main():
         )
 
     if model_args.model_name_or_path:
-        model = FlaxAutoModelForSeq2SeqLM.from_pretrained(
+        model = FlaxT5ForConditionalGeneration.from_pretrained(
             model_args.model_name_or_path, config=config, seed=training_args.seed, dtype=getattr(jnp, model_args.dtype)
         )
     else:
-        model = FlaxAutoModelForSeq2SeqLM.from_config(
+        model = FlaxT5ForConditionalGeneration.from_config(
             config, seed=training_args.seed, dtype=getattr(jnp, model_args.dtype)
         )