Logits and Labels are different shapes

[SamuelTWu]

I am using DNABERT_2 for a regression task. I load the model with num_labels=1, which should give me a regression model.

#training arguements
training_args = TrainingArguments(
    output_dir="models/DNABERT",
    learning_rate=LEARNING_RATE,
    per_device_train_batch_size=BATCH_SIZE,
    per_device_eval_batch_size=BATCH_SIZE,
    num_train_epochs=EPOCHS,
    eval_steps = 500,
    evaluation_strategy="steps",
    save_strategy="steps",
    save_total_limit=2,
    load_best_model_at_end=True,
    weight_decay=0.01,
    dataloader_num_workers = 4,
    logging_steps = 10,
)

#Regression Trainer
class RegressionTrainer(Trainer):
    def compute_loss(self, model, inputs, return_outputs=False):
        labels = inputs.pop("labels")
        outputs = model(**inputs)
        logits = outputs[0][:, 0]
        loss = torch.nn.functional.mse_loss(logits, labels)
        return (loss, outputs) if return_outputs else loss

#Model
model = AutoModelForSequenceClassification.from_pretrained("DNABERT-2-117M",  num_labels=1, trust_remote_code=True)

trainer = RegressionTrainer(
    model=model,
    args=training_args,
    train_dataset=ds["train"],
    eval_dataset=ds["valid"],
    compute_metrics=compute_metrics_for_regression
)

trainer.train()

However, when I compute the MSE/MAE/R2 of the model, I get the error:

ValueError: Found input variables with inconsistent numbers of samples: [9000, 2]

I believe this is because the logits and the labels are not the same size. Here is my compute_metrics_for_regression:

def compute_metrics_for_regression(eval_pred):
    logits, labels = eval_pred
    
    mse = mean_squared_error(labels, logits)
    mae = mean_absolute_error(labels, logits)
    r2 = r2_score(labels, logits)
    single_squared_errors = ((logits - labels).flatten()**2).tolist()
    
    return {"mse": mse, "mae": mae, "r2": r2}

How can I make the logits and labels the same size for evaluation?