Issue with mismatched tensor sizes during training with DeepSpeed

I'm currently training a model using** Hugging Face**'s Trainer with DeepSpeed integration, and I'm encountering an error related to mismatched tensor sizes. Specifically, I am getting the following error:

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

I hope someone can fix it and share your version please!😊😊😊

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.