Update test_modeling_flax_common.py

hyenal · May 14, 2024 · e290c58 · e290c58
1 parent d3af86f
commit e290c58
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tests/test_modeling_flax_common.py b/tests/test_modeling_flax_common.py
@@ -372,7 +372,7 @@ def test_equivalence_flax_to_pt(self):
                 with tempfile.TemporaryDirectory() as tmpdirname:
                     fx_model.save_pretrained(tmpdirname)
                     pt_model_loaded = pt_model_class.from_pretrained(
-                        tmpdirname, from_flax=True, attn_implementation="eager"
+                        tmpdirname, from_flax=True, attn_implementation=fx_model.config._attn_implementation
                     )
 
                 # send pytorch model to the correct device