diff --git a/tests/test_export_to_hf.py b/tests/test_export_to_hf.py
index 31d44d34f..2e86900c6 100644
--- a/tests/test_export_to_hf.py
+++ b/tests/test_export_to_hf.py
@@ -41,13 +41,7 @@ def test_export_lm_to_hf():
             config = export_lm_to_hf.ConvertLmConfig(
                 checkpoint_path=f"{tmpdir}/ckpt",
                 output_dir=f"{tmpdir}/output",
-                model=export_lm_to_hf.Gpt2Config(
-                    num_layers=2,
-                    num_heads=2,
-                    seq_len=64,
-                    use_flash_attention=True,
-                    hidden_dim=32,
-                ),
+                model=model_config,
             )
             export_lm_to_hf.main(config)