Update chatNT.py

chatNT.py

@@ -1763,6 +1763,9 @@ class TorchMultiModalPerceiverResampler(nn.Module):
             concat_input_1 = torch.cat([xf_1, x], dim=1)
             concat_input_2 = torch.cat([xf_2, x], dim=1)
 
+            outs[f"concat_input_1_{layer_idx}"] = concat_input_1.clone()
+            outs[f"concat_input_2_{layer_idx}"] = concat_input_2.clone()
+
             output = layer(
                 x=x,
                 cross_attention_embeddings_1=concat_input_1,
@@ -1771,6 +1774,7 @@ class TorchMultiModalPerceiverResampler(nn.Module):
                 attention_mask_2=attention_mask_2,
             )
             x = output["embeddings"]
+            outs[f"attention_embeddings_{layer_idx}"] = output["embeddings"].clone()
 
         return x, outs
 
@@ -1784,6 +1788,7 @@ class TorchMultiModalPerceiverResampler(nn.Module):
         """
         Computes the embeddings based on the input tokens.
         """
+        outs = {}
         assert (
             input_embeddings_1.shape[-1] == self.config.embed_dim
         ), "The input embedding dim should match the model embed dim"
@@ -1798,6 +1803,8 @@ class TorchMultiModalPerceiverResampler(nn.Module):
         outs: Dict[str, torch.Tensor] = {}
         x = latent_queries
 
+        outs["latent_queries"] = x.clone()
+
         x, outs = self.apply_attention_blocks(
             x=x,
             xf_1=input_embeddings_1,
@@ -1865,13 +1872,17 @@ class TorchMultiModalPerceiverResamplerProjection(nn.Module):
             english_token_ids, self.config.resampled_length, self.english_pad_token_id
         )
 
-        projected_embeddings = self.perceiver_resampler(
+        projected_embeddings, new_outs = self.perceiver_resampler(
             input_embeddings_1=projected_bio_embeddings,
             attention_mask_1=bio_attention_mask,
             input_embeddings_2=english_embeddings,
             attention_mask_2=english_attention_mask,
-        )["embeddings"]
+        )
+        projected_embeddings = projected_embeddings["embeddings"]
 
+        for key in new_outs.keys():
+            outs[f"{key}_perceiver"] = new_outs[key]
+            
         return projected_embeddings, outs