Update chatNT.py

chatNT.py

@@ -1763,8 +1763,8 @@ 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()
+            #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,
@@ -1774,7 +1774,7 @@ class TorchMultiModalPerceiverResampler(nn.Module):
                 attention_mask_2=attention_mask_2,
             )
             x = output["embeddings"]
-            outs[f"attention_embeddings_{layer_idx}"] = output["embeddings"].clone()
+            #outs[f"attention_embeddings_{layer_idx}"] = output["embeddings"].clone()
 
         return x, outs
 
@@ -1789,6 +1789,9 @@ class TorchMultiModalPerceiverResampler(nn.Module):
         Computes the embeddings based on the input tokens.
         """
         new_outs = {}
+        new_outs["input_embeddings_1"] = input_embeddings_1.clone()
+        new_outs["input_embeddings_2"] = input_embeddings_2.clone()
+        
         assert (
             input_embeddings_1.shape[-1] == self.config.embed_dim
         ), "The input embedding dim should match the model embed dim"
@@ -1884,7 +1887,7 @@ class TorchMultiModalPerceiverResamplerProjection(nn.Module):
         projected_embeddings = projected_embeddings["embeddings"]
 
         for key in new_outs.keys():
-            outs[f"{key}_perceiver"] = new_outs[key]
+            outs[f"PERCEIVER_{key}"] = new_outs[key]
             
         return projected_embeddings, outs