Made model Batch ready

Project_Model/Libs/Transformer/Classes/Decoder.py

@@ -35,22 +35,28 @@ class Decoder(nn.Module):
         )
         self.__layer_norm_3 = nn.LayerNorm(embedding_dimension)
 
-
-
-    def forward(self, x, k_x, v_x, padding_mask = None): #-> list[torch.Tensor]:  # k_x = v_x . While x_q = x
+    def forward(
+        self,
+        args: tuple[
+            torch.Tensor,
+            torch.Tensor,
+            torch.Tensor,
+            torch.Tensor
+        ]
+    ):  # -> list[torch.Tensor]:  # k_x = v_x . While x_q = x
+        # WARNING: args is needed to have sequential
+        x, k_x, v_x, padding_mask = args
 
         # build of attention mask
         attention_mask = get_causal_attention_mask(x.size(1))
 
         # 1) Masked Attention
         MASKED_ATTENTION = self.__masked_attention(
-            x, x, x, key_padding_mask=padding_mask, attn_mask=attention_mask
+            x, x, x, key_padding_mask=padding_mask, attention_mask=attention_mask
         )
 
         # 2) Dropout
-        DROPPED_MASKED_ATTENTION = self.__dropout(
-            MASKED_ATTENTION
-        )
+        DROPPED_MASKED_ATTENTION = self.__dropout(MASKED_ATTENTION)
         del MASKED_ATTENTION
 
         # 3) Residual Connection
@@ -61,7 +67,9 @@ class Decoder(nn.Module):
         x = self.__layer_norm_1(x)
 
         # 5) Encoder–decoder (cross) attention
-        CROSS_ATTENTION = self.__cross_attention(x, k_x, v_x, key_padding_mask=padding_mask)
+        CROSS_ATTENTION = self.__cross_attention(
+            x, k_x, v_x, key_padding_mask=padding_mask
+        )
 
         # 6) Dropout
         DROPPED_CROSS_ATTENTION = self.__dropout(CROSS_ATTENTION)
@@ -88,7 +96,7 @@ class Decoder(nn.Module):
         # 12) Layer Normalization
         x = self.__layer_norm_3(x)
 
-        return x, k_x, v_x, padding_mask
+        return (x, k_x, v_x, padding_mask)
 
 
 # use eval to disable dropout ecc