Fix test

Author: balvisio

bionemo-recipes/recipes/evo2_megatron/tests/bionemo/evo2/models/test_evo2_lora_1.py

@@ -517,19 +517,22 @@ def test_lora_checkpoint_excludes_frozen_embeddings(self, tmp_path: Path, base_c
         assert len(adapter_keys) > 0, "Checkpoint should still contain LoRA adapter keys."
 
     @pytest.mark.parametrize(
-        "skip_module, expected_key_substr, lora_targets",
+        "skip_freeze, expected_key_substr, lora_targets",
         [
-            ("word_embeddings", "word_embeddings", None),
-            ("final_norm", "final_norm", None),
-            ("dense", "mixer.dense.", None),
-            ("linear_fc2", "mlp.linear_fc2.", ["dense_projection", "linear_qkv", "linear_proj", "linear_fc1"]),
+            # word_embeddings and output_layer share a weight tensor when tying is enabled;
+            # both must appear in skip_freeze to satisfy the symmetry contract.
+            (["word_embeddings", "output_layer"], "word_embeddings", None),
+            (["final_norm"], "final_norm", None),
+            (["dense"], "mixer.dense.", None),
+            (["linear_fc2"], "mlp.linear_fc2.", ["dense_projection", "linear_qkv", "linear_proj", "linear_fc1"]),
         ],
+        ids=["word_embeddings", "final_norm", "dense", "linear_fc2"],
     )
     def test_lora_skip_freeze_saves_and_trains_module(
         self,
         tmp_path: Path,
         base_ckpt: Path,
-        skip_module: str,
+        skip_freeze: list[str],
         expected_key_substr: str,
         lora_targets: list[str] | None,
     ):
@@ -538,12 +541,12 @@ def test_lora_skip_freeze_saves_and_trains_module(
 
         from bionemo.evo2.models.evo2_provider import hyena_forward_step
 
-        lora_dir = tmp_path / f"lora_{skip_module}"
+        lora_dir = tmp_path / f"lora_{skip_freeze[0]}"
         cfg = _build_pretrain_config(
             lora_dir,
             train_iters=1,
             lora=True,
-            skip_freeze=[skip_module],
+            skip_freeze=skip_freeze,
             lora_target_modules=lora_targets,
             pretrained_ckpt_dir=str(base_ckpt),
         )