The size of tensor a (882) must match the size of tensor b (568) at non-singleton dimension 1

[QiyaoWei]

I am using quite a standard pipeline to train reward modelling with an implicit preference dataset, but I run into the issue of tensor dimension mismatch. May I ask what might be the issue here, and what debugging steps I can take to resolve this issue?

import torch
from datasets import load_dataset
from trl import RewardTrainer, RewardConfig
from transformers import AutoModelForCausalLM, AutoTokenizer
torch.set_default_device('cuda')
model = AutoModelForCausalLM.from_pretrained("gemma3", attn_implementation='eager')
tokenizer = AutoTokenizer.from_pretrained("gemma3")

# load training data, and process it so it becomes an implicit preference dataset ("chosen" and "rejected")
train_dataset = load_dataset("json", data_files="custom_training_data.json", split="train")
def prefix_with_input(example):
    example['chosen'] = example['input'] + " " + example['chosen']
    example['rejected'] = example['input'] + " " + example['rejected'][0]
    return example
train_dataset = train_dataset.map(prefix_with_input)
train_dataset = train_dataset.remove_columns(["input"])

training_args = RewardConfig()
tokenizer.pad_token = tokenizer.eos_token
training_args.dataloader_pin_memory=False
training_args.per_device_train_batch_size = 1

trainer = RewardTrainer(
    model=model,
    args=training_args,
    processing_class=tokenizer,
    train_dataset=train_dataset
)
trainer.train()

Error message below:

The size of tensor a (882) must match the size of tensor b (568) at non-singleton dimension 1
  File "train.py", line 109, in <module>
    trainer.train()
RuntimeError: The size of tensor a (882) must match the size of tensor b (568) at non-singleton dimension 1

[QiyaoWei]

Closing this issue now, as it is resolved by the forum. It would be great if the answer on the forum could just pop up within a simple Google search, because tensor dimension mismatch is quite a common problem, and the debugging process would have been much easier if I had seen a response like the one of the forum (I spent a lot of time looking at unrelated solutions). Keen to hear about how to publicize these knowledge points to make debugging easier for everyone!