diff --git a/hyperlora/modeling_utils.py b/hyperlora/modeling_utils.py
index e4f6d20..9f94847 100644
--- a/hyperlora/modeling_utils.py
+++ b/hyperlora/modeling_utils.py
@@ -19,7 +19,7 @@ from pooling import POOL_FN, get_pooling_fn
 from torch import Tensor, einsum, nn
 from transformers import PerceiverConfig, PerceiverModel, PreTrainedModel
 from transformers.models.perceiver.modeling_perceiver import (
-    PerceiverFourierPositionEncoding,
+    PerceiverBasicDecoder,
 )
 from transformers.modeling_outputs import ModelOutput
 from utils import get_lora_module_names, get_num_layers, get_peft_in_out_features
@@ -109,50 +109,43 @@ def get_hypernet_config(
 
 
 class Perceiver(nn.Module):
-    # TODO: we could expand the latent query size to be much bigger
-    # and then do cross-attn with an output query with size [n_layers x n_modules]
-    # https://huggingface.co/docs/transformers/en/model_doc/perceiver#transformers.models.perceiver.modeling_perceiver.PerceiverBasicDecoder
     """perceiver w/ bottleneck size = n_modules * n_layers"""
 
     def __init__(
         self, feature_size, output_size, num_layers, num_modules, *args, **kwargs
     ):
         super().__init__()
-        # num_bands = 256
-        # self.positional_encoding = PerceiverFourierPositionEncoding(
-        #     num_bands=num_bands, max_resolution=(128_000,), concat_pos=False,
-        # )
-        # self.pos_proj = nn.Linear(num_bands * 2, num_bands)
-        # # pos_emb_size = self.positional_encoding.output_size()
         self.num_layers = num_layers
         self.num_modules = num_modules
         self.config = PerceiverConfig(
             d_model=feature_size,  # + num_bands
-            num_latents=num_layers * num_modules,
+            num_latents=num_layers * num_modules * 8,
             d_latents=output_size,
             attention_probs_dropout_prob=0.0,
-            num_blocks=8,
-            num_self_attends_per_block=3,
+            num_blocks=1,
+            num_self_attends_per_block=26,
             **kwargs,
         )
-        # TODO: could do something more complex e.g., decoder query, etc.
-        self.perceiver = PerceiverModel(self.config)
+        decoder = PerceiverBasicDecoder(
+            self.config,
+            output_num_channels=output_size,
+            output_index_dims=num_layers * num_modules,
+            num_channels=output_size,
+            final_project=False,
+            trainable_position_encoding_kwargs=dict(
+                num_channels=output_size,
+                index_dims=num_layers * num_modules,
+            ),
+        )
+
+        self.perceiver = PerceiverModel(self.config, decoder=decoder)
 
     def forward(
         self,
         ctx_features: Float[Tensor, "bs seq_len feature_dim"],
         ctx_attn_mask: Optional[Integer[Tensor, "bs seq_len"]] = None,
     ):
-        # bs, seq_len = ctx_features.shape[:2]
-        # pos_embs = self.positional_encoding((seq_len,),
-        #                                     bs,
-        #                                     device=ctx_features.device,
-        #                                     dtype=ctx_features.dtype)
-        # pos_embs = self.pos_proj(pos_embs)
-        # pos_embs = repeat(pos_embs, "seq_len d -> bs seq_len d", bs=bs)
-        # x = torch.cat([ctx_features, pos_embs], dim=-1)
-        x = ctx_features
-        x = self.perceiver(x, ctx_attn_mask).last_hidden_state
+        x = self.perceiver(ctx_features, ctx_attn_mask).logits  # .last_hidden_state
         x = rearrange(
             x,
             "bs (n_layers n_modules) d -> bs n_layers n_modules d",
diff --git a/hyperlora/training_utils.py b/hyperlora/training_utils.py
index 4e4313e..834bfbd 100644
--- a/hyperlora/training_utils.py
+++ b/hyperlora/training_utils.py
@@ -44,18 +44,21 @@ def decode_test_result(test_dataset, test_result, tokenizer):
         if "labels" in sample:
             start_idx = np.argmax(sample["labels"] != -100)
             label_toks = sample["labels"][start_idx:]
-            # labels are padded with -100, so we need to replace them with the pad token id
+            # labels are padded with -100, so we need to
+            # replace them with the pad token id
             label_toks = np.where(label_toks == -100, tokenizer.pad_token_id, label_toks)
             label_text = tokenizer.decode(label_toks, skip_special_tokens=True)
             d["label"] = label_text
 
-        # HACK: remove the label part
+        # remove the label part
         input_toks = sample["input_ids"][:start_idx]
         gen_toks = pred_toks[len(input_toks) :]
         gen_toks = np.where(gen_toks == -100, tokenizer.pad_token_id, gen_toks)
 
         d["input"] = tokenizer.decode(input_toks, skip_special_tokens=True)
         d["generated"] = tokenizer.decode(gen_toks, skip_special_tokens=True)
+        if "ctx_ids" in sample:
+            d["context"] = tokenizer.decode(sample["ctx_ids"], skip_special_tokens=True)
         out.append(d)
 
     return out