diff --git a/src/ctx_to_lora/configs.py b/src/ctx_to_lora/configs.py
index f9d54a2..f983e01 100644
--- a/src/ctx_to_lora/configs.py
+++ b/src/ctx_to_lora/configs.py
@@ -496,9 +496,15 @@ class AggregatorArguments:
     #     default=8,
     #     metadata={"help": "Number of blocks for Perceiver."},
     # )
+
+    # misnomer
     num_self_attends_per_block: int = field(
         default=6,
-        metadata={"help": "Number of self-attends per block for Perceiver."},
+        metadata={"help": "Number of attn blocks for Perceiver."},
+    )
+    n_cross_attn_layers: int = field(
+        default=6,
+        metadata={"help": "Number of cross-attention layers for Perceiver. "},
     )
     # self_attention_widening_factor: int = field(
     #     default=4,
diff --git a/src/ctx_to_lora/modeling/aggregator.py b/src/ctx_to_lora/modeling/aggregator.py
index 0cb744f..8f0086f 100644
--- a/src/ctx_to_lora/modeling/aggregator.py
+++ b/src/ctx_to_lora/modeling/aggregator.py
@@ -42,6 +42,7 @@ class AggregatorConfig:
 
     # perceiver
     concat_latents_context: bool
+    n_cross_attn_layers: int  # number of cross-attention layers
     num_self_attends_per_block: int  # = 16
     decoder_depth: int  # 1 = only cross-attention
     num_latent_factor: int = 8
@@ -115,6 +116,7 @@ class Perceiver(nn.Module):
             input_size=feature_size,
             # the first layer is xattn
             resampler_depth=kwargs["num_self_attends_per_block"] + 1,
+            n_cross_attn_layers=kwargs["n_cross_attn_layers"],
             # resampler_n_latents=n_output_queries * num_latent_factor,
             resampler_n_latents=n_base_queries,
             intermediate_size_factor=4,
diff --git a/src/ctx_to_lora/modeling/idefics2.py b/src/ctx_to_lora/modeling/idefics2.py
index 0a53f8c..cf14a53 100644
--- a/src/ctx_to_lora/modeling/idefics2.py
+++ b/src/ctx_to_lora/modeling/idefics2.py
@@ -80,12 +80,14 @@ class Idefics2PerceiverConfig(PretrainedConfig):
         num_key_value_heads=4,
         attention_dropout=0.0,
         is_cross_attn=True,
+        n_cross_attn_layers=1,
         concat_latents_context=False,
         **kwargs,
     ):
         # for mlp
         self.concat_latents_context = concat_latents_context
         self.is_cross_attn = is_cross_attn
+        self.n_cross_attn_layers = n_cross_attn_layers
         self.input_size = input_size
         self.intermediate_size_factor = intermediate_size_factor
         # for perceiver
@@ -650,9 +652,9 @@ class Idefics2PerceiverResampler(Idefics2PreTrainedModel):
         self.hidden_size = config.hidden_size
         self.hidden_act = config.hidden_act
         self.n_latents = config.resampler_n_latents
-        self.n_x_attn_layers = 1
+        self.n_cross_attn_layers = config.n_cross_attn_layers
         self.depth = config.resampler_depth
-        self.n_self_attn_layers = self.depth - self.n_x_attn_layers
+        self.n_self_attn_layers = self.depth - self.n_cross_attn_layers
         self.rms_norm_eps = config.rms_norm_eps
         self.concat_latents_context = config.concat_latents_context
 
@@ -665,7 +667,7 @@ class Idefics2PerceiverResampler(Idefics2PreTrainedModel):
             x_attn_config.is_cross_attn = True
             self.layers = [
                 Idefics2PerceiverLayer(x_attn_config, idx)
-                for idx in range(self.n_x_attn_layers)
+                for idx in range(self.n_cross_attn_layers)
             ]
 
             config.is_cross_attn = False
@@ -811,7 +813,7 @@ class Idefics2PerceiverResampler(Idefics2PreTrainedModel):
                 )
                 compressed_context = layer_outputs[0]
         else:
-            for encoder in self.layers[: self.n_x_attn_layers]:
+            for encoder in self.layers[: self.n_cross_attn_layers]:
                 # print("Using cross-attention for resampler")
                 # print(f"position_ids: {position_ids}")
                 # print(f"cu_seq_lens_q: {cu_seq_lens_q}")
@@ -838,7 +840,7 @@ class Idefics2PerceiverResampler(Idefics2PreTrainedModel):
                 self.n_latents, device=position_ids.device, dtype=torch.int32
             ).repeat(1, bsz)
 
-            for processor in self.layers[self.n_x_attn_layers :]:
+            for processor in self.layers[self.n_cross_attn_layers :]:
                 cu_seq_lens_k = cu_seq_lens_q
                 max_length_k = max_length_q
                 # print(f"Using self-attention for resampler")