fixes and cleanup

2024-12-28 01:10:56 +00:00
parent e5fa842ff8
commit 332ce0ae85
6 changed files with 83 additions and 56 deletions
--- a/tests/e2e/integrations/convert_diff_transformer/test_convert_diff_transformer.py
+++ b/tests/e2e/integrations/convert_diff_transformer/test_convert_diff_transformer.py
@@ -130,6 +130,9 @@ def test_conversion_cli_repoduce_attentions(
 )
 def test_conversion_cli_split_heads(tmp_path: Path, base_config, attention: str):
    output_dir = tmp_path / "converted"
+
+    # Smallest model with an even number of attention heads
+    base_config["base_model"] = "HuggingFaceTB/SmolLM2-1.7B"
    base_config["output_dir"] = str(output_dir)
    base_config[attention] = True