Made model Batch ready

2025-10-07 16:37:20 +02:00
parent 109ad9f36b
commit fdece42462
4 changed files with 47 additions and 17 deletions
--- a/Project_Model/Libs/Transformer/Classes/TorchMultiHeadAttention.py
+++ b/Project_Model/Libs/Transformer/Classes/TorchMultiHeadAttention.py
@@ -8,17 +8,16 @@ class TorchMultiHeadAttention(nn.Module):
        self,
        embedding_dimension: int,
        number_of_attention_heads: int,
-        dropout: float = 0.0,
+        dropout: float = 0.0
    ):
        super().__init__()
-        self.attention = nn.MultiheadAttention(
+        self.attention = torch.nn.MultiheadAttention(
            embedding_dimension,
-            number_of_attention_heads,
+            num_heads=number_of_attention_heads,
            dropout=dropout,
            batch_first=True,
        )

-
    def forward(
        self,
        x_q: torch.Tensor,