[fix] Stack tensor-formatted labels in the CrossEncoder data collator - #4257

CrossEncoderTrainer fails on the first training step when the label column is formatted as tensors, for example after dataset.with_format("torch"), or after load_from_disk on a dataset that was saved with that format. #4231 fixed this in BaseDataCollator, but CrossEncoderDataCollator overrides __call__ with its own label handling, so CrossEncoder training still has the problem.

A 0-dim tensor counts as a collections.abc.Collection, so the collator treats each scalar label as a list of labels and returns a Python list of 0-dim tensors instead of one tensor. The pointwise losses then fail, including both losses that CrossEncoderTrainer uses by default:

Minimal reproducer:

from datasets import Dataset

from sentence_transformers.cross_encoder.data_collator import CrossEncoderDataCollator

dataset = Dataset.from_dict({"query": ["q1", "q2"], "document": ["d1", "d2"], "label": [1.0, 0.0]})
dataset = dataset.with_format("torch", columns=["label"], output_all_columns=True)
collator = CrossEncoderDataCollator(preprocess_fn=None)
print(collator([dataset[0], dataset[1]])["label"])
# main: [tensor(1.), tensor(0.)], plus a UserWarning about copy-constructing from a tensor
# this PR: tensor([1., 0.])

I found this through code analysis while checking which collators #4231 covers, not in actual use. I then trained each CrossEncoder loss for two steps on a small local model, with an unformatted and a torch-formatted dataset: the three losses above fail on main and pass with this change. MarginMSELoss and the listwise losses already accept a list of labels, so they ran, but every tensor row went through torch.tensor(tensor), which copies it and emits a UserWarning.

Change

When the labels are tensors, stack 0-dim ones into one tensor, and keep larger ones as the per-sample list that the listwise losses expect, without copying them. Python-valued labels take the same path as before.

Tests

• tests/cross_encoder/test_data_collator.py (new): torch-formatted float and integer labels collate to the same tensor as the unformatted dataset, and torch-formatted label lists stay one tensor per sample.
• tests/cross_encoder/test_trainer.py::test_trainer_tensor_formatted_labels: one CrossEncoderTrainer step with the default loss for num_labels=1 and num_labels=3 on a torch-formatted dataset. The model is a small randomly initialized BertForSequenceClassification built in the test, so nothing is downloaded. The step completes, and the weights change and stay finite.

HF_HUB_OFFLINE=1 python -m pytest tests/cross_encoder/test_data_collator.py tests/cross_encoder/test_trainer.py::test_trainer_tensor_formatted_labels tests/base/test_data_collator.py -q

ruff check, ruff format --check and typos pass on the changed files. I had no Hub access, so the CrossEncoder tests that download models did not run: in tests/cross_encoder and tests/base/test_data_collator.py, the same tests fail on main and on this branch, all with download errors. GPU training and the minimum-versions setup were not run.

AI assistance

This PR was prepared with Claude Code, which found the issue and wrote the fix, the tests and this description.

🤖 Generated with Claude Code