Lightning-AI/pytorch-lightning · error · TypeError

Unexpected error, the wrapped model should be a LightningMod

Error message

Unexpected error, the wrapped model should be a LightningModule, found {type(model).__name__}

What it means

_module_to_compiled.to_uncompiled handles an OptimizedModule by grabbing _orig_mod; if that inner object is not a LightningModule, it raises TypeError describing an unexpected wrapper — an invariant that should not normally occur.

Source

Thrown at src/lightning/pytorch/utilities/compile.py:79

    orig_module.predict_step = model.dynamo_ctx(orig_module.predict_step)  # type: ignore[method-assign]
    return orig_module


def to_uncompiled(model: Union["pl.LightningModule", "torch._dynamo.OptimizedModule"]) -> "pl.LightningModule":
    """Returns an instance of LightningModule without any compilation optimizations from a compiled model.

    .. warning::  This is an :ref:`experimental <versioning:Experimental API>` feature.

    This takes either a ``torch._dynamo.OptimizedModule`` returned by ``torch.compile()`` or a ``LightningModule``
    returned by ``from_compiled``.

    Note: this method will in-place modify the ``LightningModule`` that is passed in.

    """
    if isinstance(model, OptimizedModule):
        original = model._orig_mod
        if not isinstance(original, pl.LightningModule):
            raise TypeError(
                f"Unexpected error, the wrapped model should be a LightningModule, found {type(model).__name__}"
            )

    elif isinstance(model, pl.LightningModule):
        if model._compiler_ctx is None:
            raise ValueError(
                "`model` is required to be a compiled LightningModule. Found a non-compiled LightningModule instead."
            )
        original = model

    else:
        raise ValueError("`model` must either be an instance of OptimizedModule or LightningModule")

    ctx = original._compiler_ctx
    if ctx is not None:
        original.forward = ctx["original_forward"]  # type: ignore[method-assign]
        original.training_step = ctx["original_training_step"]  # type: ignore[method-assign]
        original.validation_step = ctx["original_validation_step"]  # type: ignore[method-assign]

View on GitHub (pinned to 9fed5c27d2)

Solutions

  1. Only pass modules produced by torch.compile of a LightningModule
  2. Upgrade torch/lightning to compatible versions
  3. Recreate the compiled model from scratch rather than reusing stale wrappers
Defensive patterns

Strategy: try-catch

Type guard

def is_valid_optimized(m) -> bool:
    from torch._dynamo import OptimizedModule
    import lightning.pytorch as pl
    return isinstance(m, OptimizedModule) and isinstance(m._orig_mod, pl.LightningModule)

Try / catch

try:
    plain = _module_to_compiled.to_uncompiled(m)
except TypeError:
    # rebuild wrapper: recompile a proper LightningModule
    plain = my_lightning_module

Prevention

When it happens

Trigger: An OptimizedModule whose _orig_mod is a plain nn.Module passed through to_uncompiled; usually a corrupted wrap or manual construction of OptimizedModule-like objects.

Common situations: Custom dynamo wrappers, monkey-patched compile, or version mismatch between torch and lightning internals.

Related errors

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