ok we're good

stanford-crfm · Dec 3, 2024 · 074d0ec · 074d0ec
1 parent 87d7665
commit 074d0ec
Showing 1 changed file with 1 addition and 7 deletions.
diff --git 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)