huggingface/transformers · error · ValueError

return_tensors must be one of ("pt", "np"), {return_tensors=

Error message

return_tensors must be one of ("pt", "np"), {return_tensors=} not supported

What it means

Raised by DataCollatorWithFlattening.__call__ when return_tensors is neither 'pt' nor 'np'. This collator concatenates samples into one packed sequence and must materialize them as PyTorch tensors or NumPy arrays (with int64/int32 dtypes for the flash-attention keys); other frameworks' tensor types are not supported, and unlike other collators there is no 'tf' or 'jax' branch.

Source

Thrown at src/transformers/data/data_collator.py:1477

                max_length = max(max_length, len(input_ids))

        if self.return_flash_attn_kwargs:
            batch["cu_seq_lens_q"] = batch["cu_seq_lens_k"] = cu_seq_lens
            batch["max_length_q"] = batch["max_length_k"] = max_length

        # FlashAttentionKwargs and seq_idx are expected to be int32s.
        if return_tensors == "pt":
            import torch

            data_cls = torch.tensor
            dtype_64 = torch.int64
            dtype_32 = torch.int32
        elif return_tensors == "np":
            data_cls = np.array
            dtype_64 = np.int64
            dtype_32 = np.int32
        else:
            raise ValueError(f'return_tensors must be one of ("pt", "np"), {return_tensors=} not supported')

        for k, v in batch.items():
            if k in self._batch_dim_keys:
                v = [v]
            # Flash attention max_len_{q,k} are python ints
            if k not in self._py_int_keys:
                batch[k] = data_cls(v, dtype=dtype_64 if k in self._int_64_keys else dtype_32)

        return batch

View on GitHub (pinned to a597f97485)

Solutions

  1. Use return_tensors='pt' (PyTorch) or 'np' (NumPy) with this collator.
  2. For JAX/TF training, take the NumPy output (return_tensors='np') and convert with jnp.asarray / tf.convert_to_tensor downstream.
  3. If you need native TF collation, use DataCollatorWithPadding or DefaultDataCollator instead of the flattening collator.

Example fix

# before
collator = DataCollatorWithFlattening(return_tensors='tf')

# after
collator = DataCollatorWithFlattening(return_tensors='np')
batch = {k: jnp.asarray(v) for k, v in collator(features).items()}  # JAX example
Defensive patterns

Strategy: validation

Validate before calling

rt = cfg.get('return_tensors', 'pt')
if rt not in ('pt', 'np'):
    raise ValueError(f'DataCollatorWithFlattening supports only pt/np, got {rt}')
collator = DataCollatorWithFlattening(return_tensors=rt)

Type guard

def is_supported_tensors(value: str) -> bool:
    return value in ('pt', 'np')

Prevention

When it happens

Trigger: Calling DataCollatorWithFlattening(...)(features, return_tensors='tf') or 'jax', or constructing it with return_tensors='tf' (that default flows into __call__).

Common situations: Copying a TensorFlow/JAX collator setup from another pipeline into the packing collator; a Trainer configured for TF attempting to use sequence packing.

Related errors


AI-assisted analysis of huggingface/transformers@a597f97485 (2026-08-14). Data as JSON: /api/errors/b8d33b60cae14f7c. Report an issue: GitHub.