AttributeError when predicting after fine-tuning mT5ForSequenceClassification for regression

Beginners
1

Hi

I used this script to fine-tune MT5ForSequenceClassification as a regression.

I passed these args:

--model_name_or_path ${MODEL_SCRATCH} \
  --train_file ${DATA_SCRATCH}/train_stata.csv \
  --validation_file ${DATA_SCRATCH}/dev_stata.csv \
  --test_file ${DATA_SCRATCH}/test_stata.csv \
  --do_regression True \
  --metric_name rmse \
  --text_column_name "linearized_input,output" \
  --label_column_name attributable \
  --do_train \
  --do_eval \
  --do_predict \
  --max_seq_length 2048 \
  --per_device_train_batch_size 8 \
  --learning_rate 1e-4 \
  --lr_scheduler_type constant \
  --ignore_mismatched_sizes True \
  --num_train_epochs 1 \
  --output_dir ${OUTPUT_DIR} \
  --overwrite_output_dir True \
  --gradient_checkpointing True \
  --gradient_accumulation_steps 1 \
  --eval_accumulation_steps 1 \
  --text_column_delimiter "[output]" \
  --save_total_limit 1 \

It trains and evaluates fine, but when it comes time to predict, I get the error:

[INFO|trainer.py:3201] 2024-01-09 11:46:40,144 >> ***** Running Prediction *****
[INFO|trainer.py:3203] 2024-01-09 11:46:40,145 >>   Num examples = 612
[INFO|trainer.py:3206] 2024-01-09 11:46:40,145 >>   Batch size = 8
100%|██████████| 77/77 [07:44<00:00,  6.03s/it]
Traceback (most recent call last):
  File "/home/s2029717/miniconda3/envs/hf_src/lib/python3.11/site-packages/numpy/core/fromnumeric.py", line 1552, in squeeze
    squeeze = a.squeeze
              ^^^^^^^^^
AttributeError: 'tuple' object has no attribute 'squeeze'

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/home/s2029717/scripts/train-stata.py", line 763, in <module>
    main()
  File "/home/s2029717/scripts/train-stata.py", line 725, in main
    predictions = np.squeeze(predictions)
                  ^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/s2029717/miniconda3/envs/hf_src/lib/python3.11/site-packages/numpy/core/fromnumeric.py", line 1554, in squeeze
    return _wrapit(a, 'squeeze', axis=axis)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/s2029717/miniconda3/envs/hf_src/lib/python3.11/site-packages/numpy/core/fromnumeric.py", line 45, in _wrapit
    result = getattr(asarray(obj), method)(*args, **kwds)
                     ^^^^^^^^^^^^
ValueError: setting an array element with a sequence. The requested array has an inhomogeneous shape after 2 dimensions. The detected shape was (2, 612) + inhomogeneous part.

I am training on a dataset where the input is a string and the label is a float. I want to predict these floats.

Help would be appreciated!