model must be a LightningModule

[adizhol-gp]

using SwaV:
when calling trainer.fit(model):

error:
model must be a LightningModule
versions:

solo-learn                    1.0.6
pytorch-lightning             2.0.2
torch                         2.0.0
torchmetrics                  0.11.4
torchvision                   0.15.1

`model` must be a `LightningModule` or `torch._dynamo.OptimizedModule`, got `SwAV`
  File "/home/ubuntu/miniconda3/envs/gpdev/lib/python3.11/site-packages/lightning/pytorch/utilities/compile.py", line 126, in _maybe_unwrap_optimized
    raise TypeError(
  File "/home/ubuntu/miniconda3/envs/gpdev/lib/python3.11/site-packages/lightning/pytorch/trainer/trainer.py", line 529, in fit
    model = _maybe_unwrap_optimized(model)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/work/microscopy_response_prediction/train_selfsup.py", line 69, in main
    trainer.fit(model, datamodule)
  File "/home/ubuntu/miniconda3/envs/gpdev/lib/python3.11/site-packages/hydra/core/utils.py", line 186, in run_job
    ret.return_value = task_function(task_cfg)
                       ^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ubuntu/miniconda3/envs/gpdev/lib/python3.11/site-packages/hydra/core/utils.py", line 260, in return_value
    raise self._return_value
  File "/home/ubuntu/miniconda3/envs/gpdev/lib/python3.11/site-packages/hydra/_internal/hydra.py", line 132, in run
    _ = ret.return_value
        ^^^^^^^^^^^^^^^^
  File "/home/ubuntu/miniconda3/envs/gpdev/lib/python3.11/site-packages/hydra/_internal/utils.py", line 458, in <lambda>
    lambda: hydra.run(
            ^^^^^^^^^^
  File "/home/ubuntu/miniconda3/envs/gpdev/lib/python3.11/site-packages/hydra/_internal/utils.py", line 223, in run_and_report
    raise ex
  File "/home/ubuntu/miniconda3/envs/gpdev/lib/python3.11/site-packages/hydra/_internal/utils.py", line 223, in run_and_report
    raise ex
  File "/home/ubuntu/miniconda3/envs/gpdev/lib/python3.11/site-packages/hydra/_internal/utils.py", line 457, in _run_app
    run_and_report(
  File "/home/ubuntu/miniconda3/envs/gpdev/lib/python3.11/site-packages/hydra/_internal/utils.py", line 394, in _run_hydra
    _run_app(
  File "/home/ubuntu/miniconda3/envs/gpdev/lib/python3.11/site-packages/hydra/main.py", line 94, in decorated_main
    _run_hydra(
  File "/work/microscopy_response_prediction/train_selfsup.py", line 75, in <module>
    main()
  File "/home/ubuntu/miniconda3/envs/gpdev/lib/python3.11/runpy.py", line 88, in _run_code
    exec(code, run_globals)
  File "/home/ubuntu/miniconda3/envs/gpdev/lib/python3.11/runpy.py", line 198, in _run_module_as_main (Current frame)
    return _run_code(code, main_globals, None,
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
TypeError: `model` must be a `LightningModule` or `torch._dynamo.OptimizedModule`, got `SwAV`

[adizhol-gp]

update:
seems like changing the pytorch_lightning import to "import lightning.pytorch" solves the problem

[vturrisi]

I'll merge PR #370 soon enough to move from the old pytorch_lightning interface to lightning.

[vturrisi]

Merged it, everything should be fine.