TypeError: setup() got an unexpected keyword argument 'stage'

[YijiaZhao]

I use pytorch_lightning '1.8.4.post0' (pip install --upgrade pytorch_lightning ) and get this error:
TypeError: setup() got an unexpected keyword argument 'stage'

1283            if self.datamodule is not None:
1284                self._call_lightning_datamodule_hook("setup", stage=fn)
1285            self._call_callback_hooks("setup", stage=fn)
1286            import pdb;pdb.set_trace()
1287 ->         self._call_lightning_module_hook("setup", stage=fn)
1288 
1289            self.strategy.barrier("post_setup")

in pytorch_lightning/trainer/trainer.py. I saw you have fixed this #12801, but what's the version should I use? I use "pip install . " in latest github repo to update the lightning, it uses lightning instead of pytorch_lightning, however, when I replace "from pytorch_lightning.callbacks import EarlyStopping, ModelCheckpoint" with "from lightning.callbacks import EarlyStopping" , then I got ModuleNotFoundError: No module named 'lightning.callbacks.

[akihironitta]

Hi @YijiaZhao! If you'd like to use PL in master, set this env var:

PACKAGE_NAME=pytorch pip install .

[akihironitta]

Or, if you prefer the new import statement, update your script as follows:

-from pytorch_lightning.callbacks import EarlyStopping
+from lightning.pytorch.callbacks import EarlyStopping

Also, some of the items are available at the top level of lightning:

-import pytorch_lightning as pl
-trainer = pl.Trainer(...)
+import lightning as L
+trainer = L.Trainer(...)

[YijiaZhao]

Thank you for your reply. I still get this error "TypeError: setup() got an unexpected keyword argument 'stage'" when I export PACKAGE_NAME=pytorch && pip install . How to avoid it?

[akihironitta]

Can I see the setup hook in your LightningModule? It's just that Lightning is trying to internally call the hook with stage=... kwarg but the error says stage=... kwarg is unexpected.

TypeError: setup() got an unexpected keyword argument 'stage'

[drexalt]

Hey @akihironitta ! Thanks for this answer, I was just looking for how to import callbacks and noticed that the docs seem to be outdated

Thanks!