can you correct error in this code - error i am getting as    5094 elif query.dim() == 2:
   5095     # Unbatched Inputs
   5096     is_batched = False

AssertionError: For batched (3-D) `query`, expected `attn_mask` to be `None`, 2-D or 3-D but found 4-D tensor instead

# gpt_from_scratch_fixed.py

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import numpy as np
from tqdm import tqdm
import random
import math

# 2. Sample Training Data
# Extended sample training data for better batching
sample_text = """
Once upon a time, in a land far, far away, there was a kingdom where magic was real.
The people lived in harmony with mystical creatures and powerful sorcerers.
Every day was filled with adventure and wonder, as tales of bravery spread across the land.
Knights roamed the forests, dragons soared the skies, and wizards crafted spells of immense power.
The kingdom thrived under the protection of its wise ruler, King Alaric, who sought peace and prosperity for all.
But darkness loomed on the horizon as an ancient evil awakened, threatening to engulf the kingdom in shadow.
Legends spoke of a hero who would rise to confront this menace, armed with courage, wisdom, and the light of hope.
As whispers of the prophecy echoed through the villages, the stage was set for an epic battle between good and evil.
"""

# Preprocess the text
sample_text = sample_text.lower()

# 3. Data Preprocessing
# Create a character-level vocabulary
chars = sorted(list(set(sample_text)))
vocab_size = len(chars)
print(f"Vocabulary Size: {vocab_size}")

# Create mappings from characters to integers and vice versa
stoi = { ch:i for i,ch in enumerate(chars) }
itos = { i:ch for i,ch in enumerate(chars) }

# Encode the entire text as integers
encoded_text = np.array([stoi[c] for c in sample_text])

# Define hyperparameters
block_size = 30  # context window size (reduced)
batch_size = 4   # reduced batch size
epochs = 100
learning_rate = 1e-3
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 3.1 Enhanced get_batches Function
def get_batches(arr, batch_size, block_size):
    """
    Generator that yields batches of input and target data.
    """
    batch_num = 0
    # Calculate the total number of possible sequences
    num_sequences = (len(arr) - 1) // block_size
    for n in range(0, num_sequences, batch_size):
        x = []
        y = []
        for i in range(n, min(n + batch_size, num_sequences)):
            start_idx = i * block_size
            end_idx = start_idx + block_size
            x_seq = arr[start_idx:end_idx]
            y_seq = arr[start_idx + 1:end_idx + 1]
            x.append(x_seq)
            y.append(y_seq)
        # Convert lists to tensors
        x = torch.tensor(x, dtype=torch.long)
        y = torch.tensor(y, dtype=torch.long)
        yield x, y
        batch_num +=1

# 4. Define the GPT Model
class GPT(nn.Module):
    def __init__(self, vocab_size, embed_size, num_heads, num_layers, block_size, dropout=0.1):
        super(GPT, self).__init__()
        self.block_size = block_size
        self.token_emb = nn.Embedding(vocab_size, embed_size)
        self.position_emb = nn.Embedding(block_size, embed_size)
        self.layers = nn.ModuleList([
            nn.TransformerDecoderLayer(d_model=embed_size, nhead=num_heads, dim_feedforward=embed_size*4, dropout=dropout)
            for _ in range(num_layers)
        ])
        self.decoder = nn.Linear(embed_size, vocab_size)
        self.dropout = nn.Dropout(dropout)
        
        # Generate causal mask
        mask = torch.tril(torch.ones(block_size, block_size)).unsqueeze(0).unsqueeze(0)
        self.register_buffer("mask", mask)

    def forward(self, x, targets=None):
        batch_size, seq_length = x.size()
        if seq_length > self.block_size:
            raise ValueError(f"Sequence length {seq_length} exceeds block size {self.block_size}")
        
        positions = torch.arange(0, seq_length).unsqueeze(0).expand(batch_size, seq_length).to(x.device)
        x = self.token_emb(x) + self.position_emb(positions)
        x = self.dropout(x)
        
        # Prepare mask
        # Shape: [1, 1, seq_length, seq_length]
        mask = self.mask[:, :, :seq_length, :seq_length]
        
        # Transformer expects [seq_length, batch_size, embed_size]
        x = x.permute(1, 0, 2)
        
        for layer in self.layers:
            x = layer(x, x, tgt_mask=mask)
        
        x = x.permute(1, 0, 2)  # [batch_size, seq_length, embed_size]
        logits = self.decoder(x)
        
        if targets is not None:
            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1))
            return logits, loss
        return logits

# Model hyperparameters
embed_size = 128
num_heads = 4
num_layers = 2

# 5. Training the Model
def train_model(model, encoded_text, epochs, batch_size, block_size, optimizer, device):
    model.train()
    for epoch in range(1, epochs + 1):
        epoch_loss = 0
        batches = 0
        for x, y in get_batches(encoded_text, batch_size, block_size):
            x = x.to(device)
            y = y.to(device)
            
            optimizer.zero_grad()
            logits, loss = model(x, y)
            loss.backward()
            optimizer.step()
            
            epoch_loss += loss.item()
            batches +=1
        avg_loss = epoch_loss / batches
        print(f"Epoch {epoch}/{epochs}, Loss: {avg_loss:.4f}")
    print("Training completed.")

# 6. Text Generation
def generate(model, start_text, length, block_size):
    model.eval()
    generated = start_text
    input_seq = [stoi.get(c, 0) for c in start_text.lower()]
    
    for _ in range(length):
        # Prepare input
        if len(input_seq) < block_size:
            input_ids = [0]*(block_size - len(input_seq)) + input_seq
        else:
            input_ids = input_seq[-block_size:]
        
        input_tensor = torch.tensor([input_ids], dtype=torch.long).to(device)
        
        with torch.no_grad():
            logits, _ = model(input_tensor)
            logits = logits[:, -1, :]  # Last time step
            probs = F.softmax(logits, dim=-1).cpu().numpy().flatten()
        
        # Sample from the distribution
        next_id = np.random.choice(len(probs), p=probs)
        next_char = itos[next_id]
        generated += next_char
        input_seq.append(next_id)
    
    return generated

# Main Function
def main():
    # Initialize the model
    model = GPT(vocab_size, embed_size, num_heads, num_layers, block_size).to(device)
    
    # Define optimizer
    optimizer = optim.Adam(model.parameters(), lr=learning_rate)
    
    # Start training
    train_model(model, encoded_text, epochs, batch_size, block_size, optimizer, device)
    
    # Example generation
    prompt = "once upon a time, "
    generated_text = generate(model, prompt, 200, block_size)
    print("\nGenerated Text:\n")
    print(generated_text)

if __name__ == "__main__":
    main()