Getting "could not broadcast input array from shape (2,2) into shape (2,)" on autotrain

I am trying to train facebook/bart-large-mnli with our dataset. I am getting this error when autotrain runs:

ValueError: could not broadcast input array from shape (2,2) into shape (2,)

I was trying this out, so I used a small dataset. The training dataset has about 65 rows, with each class having between 4-6 rows.

I am using default set of training params for this model:

{
  "lr": 0.00005,
  "epochs": 3,
  "max_seq_length": 128,
  "batch_size": 8,
  "warmup_ratio": 0.1,
  "gradient_accumulation": 1,
  "optimizer": "adamw_torch",
  "scheduler": "linear",
  "weight_decay": 0,
  "max_grad_norm": 1,
  "seed": 42,
  "logging_steps": -1,
  "auto_find_batch_size": false,
  "mixed_precision": "fp16",
  "save_total_limit": 1,
  "save_strategy": "epoch",
  "evaluation_strategy": "epoch"
}

cc @abhishek

This looks like a rookie error, but unable to pinpoint what!

I have the same issue