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

[fabrahman]

Hi,

When calling train.fit(model) I get the following TypeError. I got a couple of more similar errors but could find solutions to them (mainly due to the newer version of pl), but I could not find any fix for following error:

  File "train.py", line 540, in <module>
    main(args)
  File "train.py", line 506, in main
    logger=logger,
  File "/Structure-Aware-BART/src/lightning_base.py", line 700, in generic_train
    trainer.fit(model)
  File "/anaconda3/envs/s-bart/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py", line 458, in fit
    self._run(model)
  File "/anaconda3/envs/s-bart/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py", line 713, in _run
    self.call_setup_hook(model)  # allow user to setup lightning_module in accelerator environment
  File "/anaconda3/envs/s-bart/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py", line 1161, in call_setup_hook
    model.setup(stage=fn)
TypeError: setup() got an unexpected keyword argument 'stage'

Any pointers would be much appreciated. Thx

[ananthsub]

Can you share your LightningModule code? Are you overriding the setup function? If so, are you overriding it with this signature? https://github.com/PyTorchLightning/pytorch-lightning/blob/03bb389b218084f365b922f2a50133a5aaaadaf0/pytorch_lightning/core/hooks.py#L395

[fabrahman]

Thanks for the reply! I am using the code here. It does not look to be overriding setup.

[Borda]

can you also share what PL version are you using?
I have check the code and the signature is def setup(self, mode): which is strange...
while you are using the repo, what the training call to reproduce?

[carmocca]

Yes, you need to change it to def setup(self, stage=None):