The documents about override backward()
I'm` trying to set retain_graph=True in loss.backward()
The document says backward function could be overridden like:
https://pytorch-lightning.readthedocs.io/en/0.7.6/introduction_guide.html#extensibility
class LitMNIST(LightningModule):  
    def backward(self, use_amp, loss, optimizer):  
        # do a custom way of backward  
        loss.backward(retain_graph=True)

But the error occurs:
An exception has occurred: TypeError backward() takes 4 positional arguments but 5 were given
I check the code, it seems the arguments for backward are
model_ref.backward(self, closure_loss, optimizer, opt_idx)

  
    
      pytorch-lightning/pytorch_lightning/trainer/training_loop.py
    
    
         Line 820
      in
      1d565e1
    
  
  
    

        
          
           model_ref.backward(self, closure_loss, optimizer, opt_idx) 
        
    
  


I think this part of the document needs to be updated.