Allow configuring optimizer_step based on the gradients.
As of now, there is no way to skip optimizer_step based on the gradients after they are calculated since the backward resides inside the closure and this closure is passed to the optimizer.step. Also if training_step returns None it still calls optimizer_step due to the closure which I think it should not if there are no pending accumulated gradients.

  
    
      pytorch-lightning/pytorch_lightning/trainer/training_loop.py
    
    
        Lines 707 to 718
      in
      fde972f
    
  
  
    

        
          
           def train_step_and_backward_closure(): 
        

        
          
               result = self.training_step_and_backward( 
        

        
          
                   split_batch, 
        

        
          
                   batch_idx, 
        

        
          
                   opt_idx, 
        

        
          
                   optimizer, 
        

        
          
                   self.trainer.hiddens 
        

        
          
               ) 
        

        
          
               return None if result is None else result.loss 
        

        
          
            
        

        
          
           # optimizer step 
        

        
          
           self.optimizer_step(optimizer, opt_idx, batch_idx, train_step_and_backward_closure) 
        
    
  



  
    
      pytorch-lightning/pytorch_lightning/core/optimizer.py
    
    
        Lines 127 to 136
      in
      fde972f
    
  
  
    

        
          
           if trainer.on_tpu: 
        

        
          
               with trainer.profiler.profile(profiler_name): 
        

        
          
                   xm.optimizer_step(optimizer, optimizer_args={'closure': closure, **kwargs}) 
        

        
          
            
        

        
          
           elif trainer.amp_backend is not None: 
        

        
          
               trainer.precision_connector.backend.optimizer_step(trainer, optimizer, closure) 
        

        
          
            
        

        
          
           else: 
        

        
          
               with trainer.profiler.profile(profiler_name): 
        

        
          
                   optimizer.step(closure=closure, *args, **kwargs) 
        
    
  


Alternate option

Use manual optimization
zero_grad after backward the gradients and let optimizer_step happen which will have no effect.

Expected behavior
Would like to have this support with automatic optimization with no hacky way (2nd point above)