I’m training a model using Hugging Face’s Trainer with DeepSpeed integration, and I’m encountering an error related to mismatched tensor sizes:

The size of tensor a (50) must match the size of tensor b (3) at non-singleton dimension2

How to fix it?

➡️my data

➡️my code

  def training(data, Num, tokenizer_name="bert-base-uncased", max_length=50):
    tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)
    subset_key = list(data)[0]
    subset_data = data[subset_key][Num]
    inner_key = list(subset_data.keys())[Num]
    key_list = list(subset_data[inner_key][0].keys())
    results = [f"{key_list[0]}: {item[key_list[0]]}, {key_list[1]}: {item[key_list[1]]}"
               for item in subset_data[inner_key]]
    tokenized_results = tokenizer(results, padding="max_length", truncation=True, max_length=max_length)
    return tokenized_results['input_ids'] , tokenized_results['attention_mask']
train_dataset = training(dataset, 0, tokenizer_name="bert-base-uncased", max_length=30)

Organize and extract files ⬆️

    from transformers import Trainer, TrainingArguments
import torch

# DeepSpeed
deepspeed_config = {
    "zero_optimization": {
        "stage": 2,
        "offload_optimizer": {"device": "cpu"},
        "offload_param": {"device": "cpu"}
    },
    "train_micro_batch_size_per_gpu": 3,
    "gradient_accumulation_steps": 1,
    "fp16": {"enabled": "auto"}
}
training_args = TrainingArguments(
    output_dir="./results",
    evaluation_strategy="epoch",
    per_device_train_batch_size=3,
    per_device_eval_batch_size=3,
    num_train_epochs=3,
    logging_dir="./logs",
    logging_steps=10,
    save_steps=500,
    save_total_limit=3,
    remove_unused_columns=False,
    fp16=False,
    deepspeed=deepspeed_config
)
trainer = Trainer(
    model=model_1,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=dataset["train"],
)

device = torch.device("cpu")
model_1.to(device)

trainer.train()

main train code ⬆️

I’ve checked that the input_ids and labels have the same shape.
I’ve verified the batch size in both the Trainer configuration and the DeepSpeed config.
I’ve also ensured that the model is correctly placed on the device (cuda or cpu).

maybe, i’m not sure.