When i run this code: import torch
from torch import nn
import transformers
import json
from sklearn.model_selection import train_test_split
from torch.nn import TransformerDecoder, TransformerDecoderLayer


# Load the JSON dataset
with open('C:/Users/Dell-PC/Desktop/Projets/Datasets/summarization_dataset.json', 'r') as f:
    dataset = json.load(f)

# Split the dataset into training, validation, and testing sets
train_data, val_data = train_test_split(dataset, test_size=0.2, random_state=42)
val_data, test_data = train_test_split(val_data, test_size=0.5, random_state=42)

# Tokenize the text into words
all_words = [word for text in dataset for word in text['doc'].split()] + \
            [word for text in dataset for word in text['summary'].split()]
vocab = set(all_words)
vocab_index = {word: idx for idx, word in enumerate(vocab)}

# Adjust the preprocess_data function accordingly
def preprocess_data(data):
    input_seqs = [row['doc'].split() for row in data] # Split into words here
    target_seqs = [row['summary'].split() for row in data]
    input_tensors = [torch.tensor([vocab_index[word] for word in input_seq]) for input_seq in input_seqs]
    target_tensors = [torch.tensor([vocab_index[word] for word in target_seq]) for target_seq in target_seqs]
    return input_tensors, target_tensors

def generate_square_subsequent_mask(sz):
    mask = (torch.triu(torch.ones(sz, sz)) == 1).transpose(0, 1)
    mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))
    return mask

class DecoderOnlyTransformer(nn.Module):
    def __init__(self, n_layers, d_model, n_heads, dim_feedforward, vocab_size):
        super().__init__()
    # Embedding for input tokens
        self.token_embedding = nn.Embedding(vocab_size, d_model)
    # Positional encoding is usually needed, however not included in this minimal example
    # Transformer decoder layers
        decoder_layer = TransformerDecoderLayer(d_model=d_model, nhead=n_heads, dim_feedforward=dim_feedforward)
        self.transformer_decoder = TransformerDecoder(decoder_layer, num_layers=n_layers)
        self.out = nn.Linear(d_model, vocab_size)
        
    def forward(self, src):
    # Get token embeddings
        src = self.token_embedding(src)
    # Creating a dummy memory for the decoder since this is a decoder-only model
        memory = torch.zeros_like(src)
    # Generating a mask to prevent the model from seeing future tokens
        tgt_mask = generate_square_subsequent_mask(src.size(0)).to(src.device)
    # Forward pass through the transformer decoder
        output = self.transformer_decoder(tgt=src, memory=memory, tgt_mask=tgt_mask)
    # Pass through the final output layer
        output = self.out(output)
        return output

# Preprocess the training, validation, and testing data
train_input, train_target = preprocess_data(train_data)
val_input, val_target = preprocess_data(val_data)
test_input, test_target = preprocess_data(test_data)

# Initialize the model, loss function, and optimizer
n_layers = 6
d_model = 512
n_head = 8
dim_ff = 2048
vocab_size = len(vocab)
model = DecoderOnlyTransformer(n_layers, d_model, n_head, dim_ff, vocab_size)
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)

# Print variables for debugging
print("Train Input:", len(train_input))
print("Train Target:", len(train_target))
print("Val Input:", len(val_input))
print("Val Target:", len(val_target))
print("Test Input:", len(test_input))
print("Test Target:", len(test_target))
print("Model:", model)
print("Criterion:", criterion)
print("Optimizer:", optimizer)

# Training loop
num_epochs = 10
for epoch in range(num_epochs):
    for input_tensor, target_tensor in zip(train_input, train_target):
        optimizer.zero_grad()
        output = model(input_tensor.unsqueeze(0))

        loss = criterion(output.reshape(-1, vocab_size), target_tensor.reshape(-1))
        loss.backward()

        optimizer.step()

    # Validation
    val_loss = 0
    with torch.no_grad():
        for val_input_tensor, val_target_tensor in zip(val_input, val_target):
            val_output = model(val_input_tensor.unsqueeze(0))
            val_loss += criterion(val_output.reshape(-1, vocab_size), val_target_tensor.reshape(-1))
    val_loss /= len(val_input)

    print(f'Epoch {epoch+1}/{num_epochs}, Loss: {loss.item():.4f}, Val Loss: {val_loss.item():.4f}')   i get this error: Traceback (most recent call last):
  File "c:\Users\Dell-PC\Desktop\Projets\Mamba\Decoder-training.py", line 93, in <module>
    loss = criterion(output.reshape(-1, vocab_size), target_tensor.reshape(-1))
  File "C:\Users\Dell-PC\AppData\Local\Programs\Python\Python310\lib\site-packages\torch\nn\modules\module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "C:\Users\Dell-PC\AppData\Local\Programs\Python\Python310\lib\site-packages\torch\nn\modules\loss.py", line 1174, in forward
    return F.cross_entropy(input, target, weight=self.weight,
  File "C:\Users\Dell-PC\AppData\Local\Programs\Python\Python310\lib\site-packages\torch\nn\functional.py", line 3029, in cross_entropy
    return torch._C._nn.cross_entropy_loss(input, target, weight, _Reduction.get_enum(reduction), ignore_index, label_smoothing)
ValueError: Expected input batch_size (1355) to match target batch_size (37).