From 3e34cd00a6e658877c36fb1419b6810da67ea4a9 Mon Sep 17 00:00:00 2001
From: meilame-tayebjee <meilame.tayebjee@polytechnique.edu>
Date: Mon, 9 Feb 2026 15:52:07 +0000
Subject: [PATCH] fix: adapt tests

---
 tests/test_pipeline.py | 9 +--------
 1 file changed, 1 insertion(+), 8 deletions(-)

diff --git a/tests/test_pipeline.py b/tests/test_pipeline.py
index dfeecd2..f5f9cab 100644
--- a/tests/test_pipeline.py
+++ b/tests/test_pipeline.py
@@ -152,14 +152,7 @@ def run_full_pipeline(
         categorical_embedding_dims=model_params["categorical_embedding_dims"],
         num_classes=model_params["num_classes"],
         attention_config=attention_config,
-        label_attention_config=(
-            LabelAttentionConfig(
-                n_head=attention_config.n_head,
-                num_classes=model_params["num_classes"],
-            )
-            if label_attention_enabled
-            else None
-        ),
+        n_heads_label_attention=attention_config.n_head,
     )
 
     # Create training config