TypeError: setup() got an unexpected keyword argument 'stage'
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
Can you share your LightningModule code? Are you overriding the setup function? If so, are you overriding it with this signature? 
  
    
      pytorch-lightning/pytorch_lightning/core/hooks.py
    
    
         Line 395
      in
      03bb389
    
  
  
    

        
          
           def setup(self, stage: Optional[str] = None) -> None: