Added attention_mask

Project_Model/Libs/Transformer/Classes/TorchMultiHeadAttention.py

@@ -1,6 +1,6 @@
 import torch
 import torch.nn as nn
-
+from typing import Optional
 
 class TorchMultiHeadAttention(nn.Module):
 
@@ -9,6 +9,7 @@ class TorchMultiHeadAttention(nn.Module):
         embedding_dimension: int,
         number_of_attention_heads: int,
         dropout: float = 0.0,
+        attention_mask: Optional[torch.Tensor] = None
     ):
         super().__init__()
         self.attention = nn.MultiheadAttention(
@@ -18,12 +19,13 @@ class TorchMultiHeadAttention(nn.Module):
             batch_first=True,
         )
 
+        self.__attention_mask = attention_mask
+
     def forward(
         self,
         x_q: torch.Tensor,
         x_k: torch.Tensor,
         x_v: torch.Tensor,
-        attention_mask=None,
         key_padding_mask=None,
     ) -> torch.Tensor:
 
@@ -32,7 +34,7 @@ class TorchMultiHeadAttention(nn.Module):
         # x * Wv -> V
 
         y, _ = self.attention.forward(
-            x_q, x_k, x_v, attn_mask=attention_mask, key_padding_mask=key_padding_mask
+            x_q, x_k, x_v, attn_mask=self.__attention_mask, key_padding_mask=key_padding_mask
         )
         return y