apache/beam · error · TypeError

Batch is not an instance of torch.Tensor

Error message

Batch {batch!r} is not an instance of torch.Tensor

What it means

PytorchTypeConstraint.type_check validates that a batched value bound to a PytorchTensor[...] type hint is actually a torch.Tensor. Beam raises this TypeError during runtime type-check when the value flowing through the PCollection is not a torch.Tensor instance.

Solutions

  1. Convert the value to a tensor before it reaches the typed PCollection: torch.as_tensor(x) or torch.from_numpy(arr).
  2. In a DoFn, return torch.stack(list_of_tensors) instead of a list/ndarray when the output hint is PytorchTensor.
  3. If values are genuinely not tensors, correct the type hint (e.g. numpy typehints) instead of PytorchTensor.
  4. As a last resort disable runtime type checking (--type_check=none), though this hides the underlying bug.

Example fix

// before
class Preprocess(beam.DoFn):
  def process(self, row):
    yield row['features'].numpy()  # ndarray, hint says PytorchTensor
// after
class Preprocess(beam.DoFn):
  def process(self, row):
    yield torch.as_tensor(row['features'])  # real torch.Tensor
Defensive patterns

Strategy: type-guard

Validate before calling

def ensure_tensor(x):
    if not isinstance(x, torch.Tensor):
        x = torch.as_tensor(x)
    return x

Type guard

def is_torch_tensor(x) -> bool:
    return isinstance(x, torch.Tensor)

Try / catch

try:
    emit(batch)
except TypeError as e:
    if 'is not an instance of torch.Tensor' in str(e):
        emit(torch.as_tensor(batch))
    else:
        raise

Prevention

When it happens

Trigger: Annotating a PCollection with PytorchTensor[torch.float32, (N, 128)] while the DoFn emits numpy arrays, Python lists, tuples, or dicts; type_check runs with Beam's default runtime type checking and the element fails isinstance(batch, torch.Tensor).

Common situations: Returning numpy arrays from a DoFn annotated as PytorchTensor, forgetting torch.from_numpy()/torch.as_tensor() after preprocessing, emitting lists from beam.BatchElements instead of stacked tensors, or type-check catching values produced before a torch conversion step.

Understand the failure class

Background: Type mismatch errors: IllegalArgumentException, TypeError and type guards across 150 open-source libraries — this error's family across 150 libraries.

Related errors


AI-assisted analysis of apache/beam@12126d8942 (2026-09-13). Data as JSON: /api/errors/7ca44d2be8eb2951. Report an issue: GitHub.

Appendix: source

Thrown at sdks/python/apache_beam/typehints/pytorch_type_compatibility.py:99

  def combine_batches(self, batches):
    return torch.cat(batches, dim=self.partition_dimension)

  def get_length(self, batch):
    return batch.size(dim=self.partition_dimension)

  def estimate_byte_size(self, batch):
    return batch.nelement() * batch.element_size()


class PytorchTypeHint():
  class PytorchTypeConstraint(typehints.TypeConstraint):
    def __init__(self, dtype, shape=()):
      self.dtype = dtype
      self.shape = shape

    def type_check(self, batch):
      if not isinstance(batch, torch.Tensor):
        raise TypeError(f"Batch {batch!r} is not an instance of torch.Tensor")
      if not batch.dtype == self.dtype:
        raise TypeError(
            f"Batch {batch!r} does not have expected dtype: {self.dtype!r}")

      for dim in range(len(self.shape)):
        if not self.shape[dim] == N and not batch.shape[dim] == self.shape[dim]:
          raise TypeError(
              f"Batch {batch!r} does not have expected shape: {self.shape!r}")

    def _consistent_with_check_(self, sub):
      # TODO Check sub against batch type, and element type
      return True

    def __key(self):
      return (self.dtype, self.shape)

    def __eq__(self, other) -> bool:
      if isinstance(other, PytorchTypeHint.PytorchTypeConstraint):

View on GitHub (pinned to 12126d8942)