Lightning-AI/pytorch-lightning · error · TypeError

Expected `torch.nn.Module` or `torch.optim.Optimizer`, got:

Error message

Expected `torch.nn.Module` or `torch.optim.Optimizer`, got: {type(obj).__name__}

What it means

Lightning's meta-tensor detection helper `_has_meta_device_parameters_or_buffers` only accepts `torch.nn.Module` or `torch.optim.Optimizer`. It inspects parameters/buffers (or optimizer param_groups) for tensors on the meta device to decide whether materialization is needed. Passing any other type (raw tensor, list, dict, custom object) to code paths like `setup`, `_validate_setup`, or `_materialize_meta_tensors` triggers this TypeError.

Source

Thrown at src/lightning/fabric/utilities/init.py:115

        else:
            uninitialized_modules.add(type(submodule).__name__)

    if uninitialized_modules:
        rank_zero_warn(
            "Parameter initialization incomplete. The following modules have parameters or buffers with uninitialized"
            " memory because they don't define a `reset_parameters()` method for re-initialization:"
            f" {', '.join(uninitialized_modules)}"
        )


def _has_meta_device_parameters_or_buffers(obj: Union[Module, Optimizer], recurse: bool = True) -> bool:
    if isinstance(obj, Optimizer):
        return any(
            t.is_meta for param_group in obj.param_groups for t in param_group["params"] if isinstance(t, Parameter)
        )
    if isinstance(obj, Module):
        return any(t.is_meta for t in itertools.chain(obj.parameters(recurse=recurse), obj.buffers(recurse=recurse)))
    raise TypeError(f"Expected `torch.nn.Module` or `torch.optim.Optimizer`, got: {type(obj).__name__}")

View on GitHub (pinned to 9fed5c27d2)

Solutions

  1. Pass only `torch.nn.Module` instances to `setup`/`setup_module` and `torch.optim.Optimizer` instances to `setup_optimizers`
  2. If you have multiple objects, call setup on each individually or use the tuple form `fabric.setup(model, optimizer)` supported by the API
  3. For raw tensors, move them with `.to(fabric.device)` instead of setup

Example fix

// before
model = fabric.setup(model.parameters())  # not a Module

// after
model = fabric.setup(model)  # torch.nn.Module
optimizer = fabric.setup_optimizers(optimizer)
Defensive patterns

Strategy: type-guard

Validate before calling

assert isinstance(obj, (torch.nn.Module, torch.optim.Optimizer)), type(obj)

Type guard

def is_setuppable(obj) -> bool:
    return isinstance(obj, (torch.nn.Module, torch.optim.Optimizer))

Try / catch

try:
    fabric.setup(obj)
except TypeError as e:
    if "Expected `torch.nn.Module`" in str(e):
        raise TypeError(f"setup got unsupported object {type(obj)}") from e
    raise

Prevention

When it happens

Trigger: Calling `fabric.setup(obj)`, `fabric.setup_module(obj)`, or `fabric.setup_optimizers(obj)` with something that is neither an nn.Module nor an Optimizer (e.g. a raw tensor, a tuple of models, a LightningModule where a bare attribute was passed, or a custom class).

Common situations: Passing `(model, optimizer)` unpacked incorrectly, passing dataloaders or raw tensors to setup, wrapping objects in a way Lightning can't introspect, version changes that made this check stricter.

Related errors


AI-assisted analysis of Lightning-AI/pytorch-lightning@9fed5c27d2 (2026-08-28). Data as JSON: /api/errors/dd1296a579d9edcc. Report an issue: GitHub.