Update chatNT.py

chatNT.py

@@ -1677,8 +1677,10 @@ class TorchMultiModalPerceiverResamplerBlock(nn.Module):
         attention_mask_1: Optional[torch.Tensor] = None,
         attention_mask_2: Optional[torch.Tensor] = None,
     ) -> Dict[str, torch.Tensor]:
+        outs_news = {}
         res = x
         x = self.norm_cross_attention_1(x)
+        outs_news["ATTENTION_layer0_layer_norm_cross_attention_1"] = x.clone()
 
         attn_output = self.cross_attention_1(
             query=x,
@@ -1686,21 +1688,30 @@ class TorchMultiModalPerceiverResamplerBlock(nn.Module):
             value=cross_attention_embeddings_1,
             attention_mask=attention_mask_1,
         )["embeddings"]
+        outs_news["ATTENTION_layer1_cross_attention_layer_1"] = attn_output.clone()
         x = res + attn_output
 
         res = x
         x = self.norm_cross_attention_2(x)
+        outs_news["ATTENTION_layer2_cross_attention_layer_2"] = x.clone()
         attn_output = self.cross_attention_2(
             query=x,
             key=cross_attention_embeddings_2,
             value=cross_attention_embeddings_2,
             attention_mask=attention_mask_2,
         )["embeddings"]
+        outs_news["ATTENTION_layer3_cross_attention_layer_2"] = attn_output.clone()
         x = res + attn_output
 
         x = x + self.mlp(x)
+        outs_news["ATTENTION_after_mlp"] = x.clone()
 
-        return {"embeddings": x}
+        output = {}
+        for key in outs_news.keys():
+            output[key] = outs_news[key]
+
+        output["embeddings"] = x
+        return output
 
 
 class TorchMultiModalPerceiverResampler(nn.Module):
@@ -1763,8 +1774,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"PERCEIVER_RESAMPLER_concat_input_1_{layer_idx}"] = concat_input_1.clone()
-            outs[f"PERCEIVER_RESAMPLER_concat_input_2_{layer_idx}"] = concat_input_2.clone()
+            #outs[f"PERCEIVER_RESAMPLER_concat_input_1_{layer_idx}"] = concat_input_1.clone()
+            #outs[f"PERCEIVER_RESAMPLER_concat_input_2_{layer_idx}"] = concat_input_2.clone()
 
             output = layer(
                 x=x,
@@ -1774,7 +1785,11 @@ class TorchMultiModalPerceiverResampler(nn.Module):
                 attention_mask_2=attention_mask_2,
             )
             x = output["embeddings"]
-            outs[f"PERCEIVER_RESAMPLER_attention_embeddings_{layer_idx}"] = output["embeddings"].clone()
+            #outs[f"PERCEIVER_RESAMPLER_attention_embeddings_{layer_idx}"] = output["embeddings"].clone()
+
+            for key in output.keys():
+                if key != "embeddings":
+                    outs[f"{key}_{layer_idx}"] = output[key].clone()
 
         return x, outs