JD

Getting the following error when running trainer.train():

Traceback (most recent call last):
  File "/smollm2/train.py", line 63, in 
    trainer.train()
  File "/smollm2/smol-env/lib/python3.12/site-packages/transformers/trainer.py", line 2238, in train
    return inner_training_loop(
           ^^^^^^^^^^^^^^^^^^^^
  File "/smollm2/smol-env/lib/python3.12/site-packages/transformers/trainer.py", line 2582, in _inner_training_loop
    tr_loss_step = self.training_step(model, inputs, num_items_in_batch)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/smollm2/smol-env/lib/python3.12/site-packages/trl/trainer/sft_trainer.py", line 904, in training_step
    return super().training_step(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/smollm2/smol-env/lib/python3.12/site-packages/transformers/trainer.py", line 3796, in training_step
    loss = self.compute_loss(model, inputs, num_items_in_batch=num_items_in_batch)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/smollm2/smol-env/lib/python3.12/site-packages/trl/trainer/sft_trainer.py", line 886, in compute_loss
    correct_predictions = (predictions == shift_labels) & mask
                           ^^^^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: The size of tensor a (2) must match the size of tensor b (4) at non-singleton dimension 0