Как создать PyTorch-модель, совместимую с Trainer из Hugging Face?

Я хочу создать пользовательский PyTorch-классификатор,
совместимый с API класса Trainer из Hugging Face.

Использую для этого простой набор данных.

import evaluate
import numpy as np

from datasets import load_dataset
from transformers import TrainingArguments, Trainer

import torch
from torch import nn

dataset = load_dataset("7wolf/y2023-e-cosmos")
dataset = dataset["train"].train_test_split(test_size=0.2)


class CustomClassifier(nn.Module):
    def __init__(self, input_dim=1, num_labels=1):
        super().__init__()
        self.sequential = nn.Linear(input_dim, num_labels)

    def forward(self, x):
        x = self.sequential(x)
        return x


input_dim = len(dataset["train"].features) - 1
num_labels = 1
model = CustomClassifier(input_dim, num_labels)


accuracy = evaluate.load("accuracy")


def compute_metrics(eval_pred):
    predictions = np.argmax(eval_pred.predictions, axis=1)
    return accuracy.compute(predictions=predictions, references=eval_pred.label_ids)


training_args = TrainingArguments(
    output_dir="nn-geometry-classification",
    evaluation_strategy="epoch",
    save_strategy="epoch",
    learning_rate=3e-5,
    per_device_train_batch_size=32,
    gradient_accumulation_steps=4,
    per_device_eval_batch_size=32,
    num_train_epochs=6,
    warmup_ratio=0.1,
    weight_decay=1e-3,
    logging_steps=10,
    load_best_model_at_end=True,
    metric_for_best_model="accuracy",
    push_to_hub=True,
    label_names=["target"],
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=dataset["train"],
    eval_dataset=dataset["test"],
    compute_metrics=compute_metrics,
)

trainer.train()

Получаю исключение:

TypeError: CustomClassifier.forward() got an unexpected keyword argument 'target'

Что я делаю не так?

Как обучить простую линейную модель при помощи Trainer?