When we run a test for this code: import torch
import torch.nn as nn
import torch.nn.functional as F
import math

class GroupedQueryAttention(nn.Module):
    def __init__(self, num_heads, hidden_size, group_size):
        super(GroupedQueryAttention, self).__init__()
        self.num_heads = num_heads
        self.hidden_size = hidden_size
        self.group_size = group_size

        self.query_linear = nn.Linear(hidden_size, hidden_size)
        self.key_linear = nn.Linear(hidden_size, hidden_size)
        self.value_linear = nn.Linear(hidden_size, hidden_size)

    def forward(self, query, key, value, attention_mask):
        print("Query shape:", query.shape)
        print("Key shape:", key.shape)
        print("Value shape:", value.shape)
        print("Attention mask shape:", attention_mask.shape)

        batch_size = 1
        seq_len = query.shape[1]
        hidden_size = self.hidden_size

        print("Batch size:", batch_size)
        print("Sequence length:", seq_len)
        print("Hidden size:", hidden_size)

        # Add hidden size dimension to query, key, and value
        query = query.unsqueeze(-1).expand(-1, -1, hidden_size)
        key = key.unsqueeze(-1).expand(-1, -1, hidden_size)
        value = value.unsqueeze(-1).expand(-1, -1, hidden_size)

        print("Query shape after adding hidden size dimension:", query.shape)
        print("Key shape after adding hidden size dimension:", key.shape)
        print("Value shape after adding hidden size dimension:", value.shape)

        # Reshape query, key, and value to group them
        query = query.view(batch_size, seq_len // self.group_size, self.group_size, hidden_size)
        key = key.view(batch_size, seq_len // self.group_size, self.group_size, hidden_size)
        value = value.view(batch_size, seq_len // self.group_size, self.group_size, hidden_size)

        print("Query shape after reshaping:", query.shape)
        print("Key shape after reshaping:", key.shape)
        print("Value shape after reshaping:", value.shape)

        # Compute attention scores
        attention_scores = torch.matmul(query, key.transpose(-1, -2)) / math.sqrt(hidden_size)
        attention_scores = attention_scores + attention_mask.unsqueeze(1).unsqueeze(1)

        print("Attention scores shape:", attention_scores.shape)

        # Compute attention weights
        attention_weights = F.softmax(attention_scores, dim=-1)

        print("Attention weights shape:", attention_weights.shape)

        # Compute output
        output = torch.matmul(attention_weights, value)
        output = output.view(batch_size, seq_len, hidden_size)

        print("Output shape:", output.shape)

        return output

class SlidingWindowAttention(nn.Module):
    def __init__(self, num_heads, hidden_size, window_size):
        super(SlidingWindowAttention, self).__init__()
        self.num_heads = num_heads
        self.hidden_size = hidden_size
        self.window_size = window_size

        self.query_linear = nn.Linear(hidden_size, hidden_size)
        self.key_linear = nn.Linear(hidden_size, hidden_size)
        self.value_linear = nn.Linear(hidden_size, hidden_size)

    def forward(self, query, key, value, attention_mask):
        batch_size, seq_len, hidden_size = query.shape
        window_size = self.window_size

        # Add extra dimension to key tensor
        key = key.unsqueeze(-1).expand(-1, -1, hidden_size)

        # Reshape query and key tensors
        query = query.view(-1, hidden_size)
        key = key.view(-1, hidden_size)

        # Compute attention scores
        attention_scores = torch.zeros(batch_size, seq_len, seq_len)
        for i in range(seq_len):
            start = max(0, i - window_size)
            end = min(seq_len, i + window_size + 1)
            query_unsqueezed = query.unsqueeze(-1)  # Add extra dimension to query tensor
            attention_scores[:, i, start:end] = torch.matmul(query_unsqueezed[:, i:i+1], key[:, start:end].transpose(-1, -2)) / math.sqrt(hidden_size)
        attention_scores = attention_scores + attention_mask.unsqueeze(1)

        # Compute attention weights
        attention_weights = F.softmax(attention_scores, dim=-1)

        # Compute output
        output = torch.matmul(attention_weights, value)
        return output

class DecoderTransformer(nn.Module):
    def __init__(self, num_layers, num_heads, hidden_size, group_size, window_size):
        super(DecoderTransformer, self).__init__()
        self.num_layers = num_layers
        self.num_heads = num_heads
        self.hidden_size = hidden_size
        self.group_size = group_size
        self.window_size = window_size

        self.layers = nn.ModuleList([DecoderTransformerLayer(num_heads, hidden_size, group_size, window_size) for _ in range(num_layers)])

    def forward(self, input_ids, attention_mask):
        output = input_ids
        for layer in self.layers:
            output = layer(output, attention_mask)
        return output

class DecoderTransformerLayer(nn.Module):
    def __init__(self, num_heads, hidden_size, group_size, window_size):
        super(DecoderTransformerLayer, self).__init__()
        self.num_heads = num_heads
        self.hidden_size = hidden_size
        self.group_size = group_size
        self.window_size = window_size

        self.self_attn = GroupedQueryAttention(num_heads, hidden_size, group_size)
        self.encoder_attn = SlidingWindowAttention(num_heads, hidden_size, window_size)
        self.feed_forward = nn.Linear(hidden_size, hidden_size)

    def forward(self, input_ids, attention_mask):
        output = self.self_attn(input_ids, input_ids, input_ids, attention_mask)
        output = output + input_ids.unsqueeze(-1).expand(-1, -1, 512)  # Add this line
        output = self.encoder_attn(output, input_ids, input_ids, attention_mask)
        output = output + input_ids.unsqueeze(-1).expand(-1, -1, 512)  # Add this line
        output = self.feed_forward(output)
        output = output + input_ids.unsqueeze(-1).expand(-1, -1, 512)  # Add this line
        return output   we get this: Query shape: torch.Size([1, 16])
Key shape: torch.Size([1, 16])
Value shape: torch.Size([1, 16])
Attention mask shape: torch.Size([1, 16])
Batch size: 1
Sequence length: 16
Hidden size: 512
Query shape after adding hidden size dimension: torch.Size([1, 16, 512])
Key shape after adding hidden size dimension: torch.Size([1, 16, 512])
Value shape after adding hidden size dimension: torch.Size([1, 16, 512])
Query shape after reshaping: torch.Size([1, 1, 16, 512])
Key shape after reshaping: torch.Size([1, 1, 16, 512])
Value shape after reshaping: torch.Size([1, 1, 16, 512])
Attention scores shape: torch.Size([1, 1, 16, 16])
Attention weights shape: torch.Size([1, 1, 16, 16])
Output shape: torch.Size([1, 16, 512])
Traceback (most recent call last):
  File "c:\Users\L14\Documents\Projets\Best-Transformer\Testing.py", line 24, in <module>
    output = model(input_ids, attention_mask)
  File "C:\Users\L14\AppData\Local\Packages\PythonSoftwareFoundation.Python.3.10_qbz5n2kfra8p0\LocalCache\local-packages\Python310\site-packages\torch\nn\modules\module.py", line 1511, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "C:\Users\L14\AppData\Local\Packages\PythonSoftwareFoundation.Python.3.10_qbz5n2kfra8p0\LocalCache\local-packages\Python310\site-packages\torch\nn\modules\module.py", line 1520, in _call_impl
    return forward_call(*args, **kwargs)
  File "c:\Users\L14\Documents\Projets\Best-Transformer\Main.py", line 120, in forward
    output = layer(output, attention_mask)
  File "C:\Users\L14\AppData\Local\Packages\PythonSoftwareFoundation.Python.3.10_qbz5n2kfra8p0\LocalCache\local-packages\Python310\site-packages\torch\nn\modules\module.py", line 1511, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "C:\Users\L14\AppData\Local\Packages\PythonSoftwareFoundation.Python.3.10_qbz5n2kfra8p0\LocalCache\local-packages\Python310\site-packages\torch\nn\modules\module.py", line 1520, in _call_impl
    return forward_call(*args, **kwargs)
  File "c:\Users\L14\Documents\Projets\Best-Transformer\Main.py", line 138, in forward
    output = self.encoder_attn(output, input_ids, input_ids, attention_mask)
  File "C:\Users\L14\AppData\Local\Packages\PythonSoftwareFoundation.Python.3.10_qbz5n2kfra8p0\LocalCache\local-packages\Python310\site-packages\torch\nn\modules\module.py", line 1511, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "C:\Users\L14\AppData\Local\Packages\PythonSoftwareFoundation.Python.3.10_qbz5n2kfra8p0\LocalCache\local-packages\Python310\site-packages\torch\nn\modules\module.py", line 1520, in _call_impl
    return forward_call(*args, **kwargs)
  File "c:\Users\L14\Documents\Projets\Best-Transformer\Main.py", line 96, in forward
    attention_scores[:, i, start:end] = torch.matmul(query_unsqueezed[:, i:i+1], key[:, start:end].transpose(-1, -2)) / math.sqrt(hidden_size)
RuntimeError: Expected size for first two dimensions of batch2 tensor to be: [16, 1] but got: [16, 16].