Update chatNT.py

chatNT.py

@@ -1788,7 +1788,7 @@ class TorchMultiModalPerceiverResampler(nn.Module):
         """
         Computes the embeddings based on the input tokens.
         """
-        outs = {}
+        new_outs = {}
         assert (
             input_embeddings_1.shape[-1] == self.config.embed_dim
         ), "The input embedding dim should match the model embed dim"
@@ -1803,7 +1803,7 @@ class TorchMultiModalPerceiverResampler(nn.Module):
         outs: Dict[str, torch.Tensor] = {}
         x = latent_queries
 
-        outs["latent_queries"] = x.clone()
+        new_outs["latent_queries"] = x.clone()
 
         x, outs = self.apply_attention_blocks(
             x=x,
@@ -1814,9 +1814,12 @@ class TorchMultiModalPerceiverResampler(nn.Module):
             attention_mask_2=attention_mask_2,
         )
 
+        for key in outs.keys():
+            new_outs[key] = outs[key].copy()
+            
         outs["embeddings"] = x
 
-        return outs
+        return outs, new_outs
 
 
 class TorchMultiModalPerceiverResamplerProjection(nn.Module):