Lightning-AI/pytorch-lightning · error · TypeError

`model` must be a `LightningModule` or `torch._dynamo.Optimi

Error message

`model` must be a `LightningModule` or `torch._dynamo.OptimizedModule`, got `{type(model).__qualname__}`

What it means

Trainer entry points (fit/validate/test/predict) call _maybe_unwrap_optimized to normalize the model. If it is neither an OptimizedModule nor a LightningModule (after a mixed-imports check), TypeError is raised naming the offending type.

Source

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

    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]
        original.test_step = ctx["original_test_step"]  # type: ignore[method-assign]
        original.predict_step = ctx["original_predict_step"]  # type: ignore[method-assign]
        original._compiler_ctx = None

    return original


def _maybe_unwrap_optimized(model: object) -> "pl.LightningModule":
    if isinstance(model, OptimizedModule):
        return from_compiled(model)
    if isinstance(model, pl.LightningModule):
        return model
    _check_mixed_imports(model)
    raise TypeError(
        f"`model` must be a `LightningModule` or `torch._dynamo.OptimizedModule`, got `{type(model).__qualname__}`"
    )


def _verify_strategy_supports_compile(model: "pl.LightningModule", strategy: Strategy) -> None:
    if model._compiler_ctx is not None:
        supported_strategies = (SingleDeviceStrategy, DDPStrategy, FSDPStrategy)
        if not isinstance(strategy, supported_strategies) or isinstance(strategy, DeepSpeedStrategy):
            supported_strategy_names = ", ".join(s.__name__ for s in supported_strategies)
            raise RuntimeError(
                f"Using a compiled model is incompatible with the current strategy: `{type(strategy).__name__}`."
                f" Only {supported_strategy_names} support compilation. Either switch to one of the supported"
                " strategies or avoid passing in compiled model."
            )

View on GitHub (pinned to 9fed5c27d2)

Solutions

  1. Make your model subclass lightning.pytorch.LightningModule (or the pytorch_lightning one matching your Trainer import)
  2. Unify all imports to a single namespace: use either `lightning.pytorch` or `pytorch_lightning`, never both

Example fix

# before
import torch.nn as nn
from pytorch_lightning import LightningModule
class Model(LightningModule): ...
from lightning.pytorch import Trainer
Trainer().fit(Model())
# after
from lightning.pytorch import LightningModule, Trainer
class Model(LightningModule): ...
Trainer().fit(Model())
Defensive patterns

Strategy: type-guard

Type guard

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

Try / catch

try:
    trainer.fit(model)
except TypeError as e:
    if "must be a" in str(e):
        raise TypeError(f"Wrap {type(model).__name__} in a LightningModule subclass") from e
    raise

Prevention

When it happens

Trigger: trainer.fit(nn.Module()) or passing any non-LightningModule model; also a LightningModule subclassed from pytorch_lightning while the Trainer is from lightning.pytorch (mixed imports).

Common situations: Forgetting to subclass LightningModule; migrating from pytorch_lightning to the lightning package but leaving some imports old.

Related errors


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