`LightningCLI` fails to recognize custom `Callback` due to mixed import styles

[tonyzamyatin]

Bug description

I implemented a custom PyTorch Lightning Callback and passed it to the Trainer via LightningCLI. However, I encountered an error where my callback was not recognized as an instance of Callback inside _validate_callbacks_list().

After debugging, I found that the issue stems from _check_mixed_imports() in lightning/pytorch/utilities/model_helpers.py, which does not correctly detect mixed imports when using lightning.pytorch vs pytorch_lightning.

What version are you seeing the problem on?

v2.5

How to reproduce the bug

import torch
from lightning.pytorch.demos.boring_classes import BoringModel as BoringLightningModule
from lightning.pytorch.demos.dummy_data import DummyDataModule as DummyDatamodule

from lightning.pytorch.cli import LightningCLI    # New import syntax
from pytorch_lightning.callbacks import Callback  # Old import syntax


class MyCallback(Callback):
    def __init__(self):
        super().__init__()
    # More code...


class MyCLI(LightningCLI):
    def add_arguments_to_parser(self, parser):
        super().add_arguments_to_parser(parser)
        parser.add_optimizer_args(torch.optim.SGD)

def cli_main():
    return MyCLI(
        BoringLightningModule,
        DummyDatamodule,
        trainer_defaults={"callbacks": [MyCallback]}
    )

if __name__ == "__main__":
    torch.set_float32_matmul_precision("medium")
    cli = cli_main()

Error messages and logs

  File "/opt/conda/lib/python3.11/site-packages/lightning/pytorch/trainer/connectors/callback_connector.py", line 228, in _validate_callbacks_list
    stateful_callbacks = [cb for cb in callbacks if is_overridden("state_dict", instance=cb)]
                         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/conda/lib/python3.11/site-packages/lightning/pytorch/utilities/model_helpers.py", line 42, in is_overridden
    raise ValueError("Expected a parent")
ValueError: Expected a parent

This suggests that the callback is not recognized as a valid subclass of Callback inside _validate_callbacks_list().

Environment

Current environment

#- PyTorch Lightning Version (e.g., 2.5.0): 2.5.0
#- PyTorch Version (e.g., 2.5): 2.6
#- Python version (e.g., 3.12): 3.11
#- OS (e.g., Linux): Linux Ubuntu
#- CUDA/cuDNN version: 12.6
#- How you installed Lightning(`conda`, `pip`, source): `pip`

More info

I traced the issue to _check_mixed_imports() in lightning/pytorch/utilities/model_helpers.py, which should detect mixed imports (e.g., pytorch_lightning vs. lightning.pytorch):

⚠️ Current _check_mixed_imports() (Bugged Behavior)

def _check_mixed_imports(instance: object) -> None:
    old, new = "pytorch_" + "lightning", "lightning." + "pytorch"
    klass = type(instance)
    module = klass.__module__

    if module.startswith(old) and __name__.startswith(new):
        pass
    elif module.startswith(new) and __name__.startswith(old):
        old, new = new, old
    else:
        return
    raise TypeError(
        f"You passed a `{old}` object ({klass.__qualname__}) to a `{new}` Trainer. "
        "Please switch to a single import style."
    )

✅ Expected Behavior: This function should catch mismatched imports and raise a TypeError.
❌ Observed Behavior: The function does not detect the mixed import case, leading to a silent error.

🔧 Proposed Fix
Modify _check_mixed_imports() to actually catch mixed imports by ensuring it checks both module and type(instance) correctly:

def _check_mixed_imports(instance: object) -> None:
    old, new = "pytorch_" + "lightning", "lightning." + "pytorch"
    klass = type(instance)
    module = klass.__module__

    if module.startswith(old) and __name__.startswith(new) or module.startswith(new) and __name__.startswith(old):
        raise TypeError(
            f"You passed a `{old}` object ({klass.__qualname__}) to a `{new}` Trainer. "
            "Please switch to a single import style."
        )

cc @lantiga @mauvilsa

[oseymour]

I'm encountering the same issue on v2.5.0.post0. Let me know if you want code to reproduce! It's basically the same as what @tonyzamyatin has already shared though.

EDIT: For people who don't want to wait for a fix, I was originally importing Callback as from pytorch_lightning.callbacks import Callback. Changed this to from lightning.pytorch.callbacks import Callback and no longer get an error.

[mauvilsa]

I agree that if someone makes the mistake of mixing imports, the error should show the You passed ... Please switch to a single import style. Otherwise it is difficult to know what to fix. Though, I would suggest to change the issue title. The problem is not related to LightningCLI. And failing to recognize custom callbacks when imports are mixed is the expected behavior. The actual problem is the confusing error message.

[mauvilsa]

As I explained in my previous comment, this is not really related to LightningCLI.

[littlebullGit]

The reported issue stems from how callbacks are passed to [LightningCLI]src/lightning/pytorch/cli.py's trainer_defaults parameter. This did not seems an issue for the mix imports.

@tonyzamyatin @mauvilsa @oseymour @SkafteNicki

Root Cause

The callbacks parameter in trainer_defaults expects instances of callbacks (e.g., [MyCallback()]), not the class itself ([MyCallback]).

Testing Results

Below is a test script to verify this behavior across different PyTorch Lightning versions (2.5.0 and 2.8.0).

# Test 1: Passing callback class - THIS WILL FAIL (expected)
try:
    cli = MyCLI(
        BoringModel,
        BoringDataModule,
        trainer_defaults={
            "callbacks": [MyCallback],  # Passing class (incorrect)
            "max_epochs": 1,
        },
        run=False
    )
    print("Test 1 PASSED: Successfully created CLI with callback class")
except Exception as e:
    print(f"Test 1 FAILED: {type(e).__name__}: {str(e)}")

# Test 2: Passing callback instance - THIS WILL PASS
try:
    cli = MyCLI(
        BoringModel,
        BoringDataModule,
        trainer_defaults={
            "callbacks": [MyCallback()],  # Passing instance (correct)
            "max_epochs": 1,
        },
        run=False
    )
    print("Test 2 PASSED: Successfully created CLI with callback instance")
except Exception as e:
    print(f"Test 2 FAILED: {type(e).__name__}: {str(e)}")

Test 1 FAILED: ValueError: Expected a parent
Test 2 PASSED: Successfully created CLI with callback instance

[SkafteNicki]

@littlebullGit thanks for your report on this issue. Makes total sense why this error is raised. I think we can close this issue now. In summary:

• Do not mix import styles in lightning
• When passing callbacks to trainer_defaults the callbacks needs to be instances of callbacks and not classes