Проблемы с обучением модели torch на google colab

я пытаюсь обучить модель в google colab и у меня выдает ошибку. Я не особо разбираюсь в обучении моделей и могу совершать глупости, помогите пожалуйста в исправлении

google colab

import torch
from datasets import load_dataset
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, Trainer, TrainingArguments, DataCollatorForSeq2Seq, AutoModelForMaskedLM

tokenizer = AutoTokenizer.from_pretrained("kazzand/ru-longformer-tiny-16384")
model = AutoModelForMaskedLM.from_pretrained("kazzand/ru-longformer-tiny-16384")

cnn_dataset = load_dataset("cnn_dailymail", '3.0.0')

def preprocess_function(examples):
    inputs = examples['article']
    targets = examples['highlights']

    # Токенизация входных данных с ограничением длины
    model_inputs = tokenizer(inputs, max_length=1024, truncation=True, padding='max_length')

    # Токенизация меток
    labels = tokenizer(targets, max_length=150, truncation=True, padding='max_length')

    # Добавляем метки в токенизированные входные данные
    model_inputs['labels'] = labels['input_ids']
    return model_inputs

# Токенизация всего датасета
tokenized_cnn = cnn_dataset.map(preprocess_function, batched=True)

data_collator = DataCollatorForSeq2Seq(tokenizer, model=model)

# Очистка кэша перед обучением
torch.cuda.empty_cache()

# Настройка параметров тренировки
training_args = TrainingArguments(
    output_dir='./results',
    evaluation_strategy="epoch",
    learning_rate=2e-5,
    per_device_train_batch_size=4,  # Уменьшение размера батча
    per_device_eval_batch_size=4,
    gradient_accumulation_steps=8,  # Увеличение шагов аккумуляции градиентов
    num_train_epochs=3,
    weight_decay=0.01,
    logging_dir='./logs',
    logging_steps=10,
    report_to="none",  # Отключение стороннего логирования
    fp16=True  # Смешанная точность для экономии памяти
)

# Включение gradient checkpointing для экономии памяти
model.gradient_checkpointing_enable()

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_cnn['train'],
    eval_dataset=tokenized_cnn['validation'],
    data_collator=data_collator  # Выравнивание размеров батчей
)

# Запуск обучения
trainer.train()

Ошибка:

ValueError                                Traceback (most recent call last)

<ipython-input-16-fe4915475509> in <cell line: 2>()
      1 # Запуск обучения
----> 2 trainer.train()

13 frames

/usr/local/lib/python3.10/dist-packages/torch/nn/functional.py in cross_entropy(input, target, weight, size_average, ignore_index, reduce, reduction, label_smoothing)
   3102     if size_average is not None or reduce is not None:
   3103         reduction = _Reduction.legacy_get_string(size_average, reduce)
-> 3104     return torch._C._nn.cross_entropy_loss(input, target, weight, _Reduction.get_enum(reduction), ignore_index, label_smoothing)
   3105 
   3106 

ValueError: Expected input batch_size (4096) to match target batch_size (600).