# Codeblock 17b
   def forward(self, features, captions, attn_mask):  #(1)
       print(f"attn_mask\t\t: {attn_mask.shape}")
       residual = captions
       print(f"captions & residual\t: {captions.shape}")
      
       #(2)
       captions, self_attn_weights = self.self_attention(query=captions,
                                                         key=captions,
                                                         value=captions,
                                                         attn_mask=attn_mask)
       print(f"after self attention\t: {captions.shape}")
       print(f"self attn weights\t: {self_attn_weights.shape}")
      
       captions = self.layer_norm_0(captions + residual)
       print(f"after norm\t\t: {captions.shape}")
      
      
       print(f"\nfeatures\t\t: {features.shape}")
       residual = captions
       print(f"captions & residual\t: {captions.shape}")
      
       #(3)
       captions, cross_attn_weights = self.cross_attention(query=captions,
                                                           key=features,
                                                           value=features)
       print(f"after cross attention\t: {captions.shape}")
       print(f"cross attn weights\t: {cross_attn_weights.shape}")
      
       captions = self.layer_norm_1(captions + residual)
       print(f"after norm\t\t: {captions.shape}")
      
       residual = captions
       print(f"\ncaptions & residual\t: {captions.shape}")
      
       captions = self.ffn(captions)  #(4)
       print(f"after ffn\t\t: {captions.shape}")
      
       captions = self.layer_norm_2(captions + residual)
       print(f"after norm\t\t: {captions.shape}")
      
       return captions