Example of gradient accumulation documentation is not correct
📚 Documentation
Example of gradient accumulation documentation is not correct :
from pytorch_lightning import Trainer
from pytorch_lightning.callbacks import GradientAccumulationScheduler

# at epoch 5 start accumulating every 2 batches
accumulator = GradientAccumulationScheduler(scheduling: {5: 2})
Trainer(accumulate_grad_batches=accumulator)
If running the code like this, I receive following error :

TypeError: Gradient accumulation supports only int and dict types


After reading the source code :

  
    
      pytorch-lightning/pytorch_lightning/trainer/training_tricks.py
    
    
        Lines 32 to 41
      in
      d856989
    
  
  
    

        
          
           def configure_accumulated_gradients(self, accumulate_grad_batches): 
        

        
          
               self.accumulate_grad_batches = None 
        

        
          
            
        

        
          
               if isinstance(accumulate_grad_batches, dict): 
        

        
          
                   self.accumulation_scheduler = GradientAccumulationScheduler(accumulate_grad_batches) 
        

        
          
               elif isinstance(accumulate_grad_batches, int): 
        

        
          
                   schedule = {1: accumulate_grad_batches} 
        

        
          
                   self.accumulation_scheduler = GradientAccumulationScheduler(schedule) 
        

        
          
               else: 
        

        
          
                   raise TypeError("Gradient accumulation supports only int and dict types") 
        
    
  


It seems the option accumulate_grad_batches can be only int or dict.