PyTorch Lightning Optimizer_Step() prevents training_step() from running
Hello,
I had an error where the training_step() was not run properly. I just found out the cause was because of the optimizer_step(). My training step immediately runs after I commented out optimizer_step().
Some other users also experienced the same error as described here: https://stackoverflow.com/questions/66756245/training-step-not-executing-in-pytorch-lightning
My question is: Now that training_step() is running, but my train_loss is explosive because of the lack of a learning rate scheduler, henceforth, what can I implement to re-enable back my learning rate scheduler?
Here's my chunk of code:
    def configure_optimizers(self):
        
        """
        AdamW Optimizer lr=0.0006

        """        
        optimizer = optim.AdamW(self.parameters(),
                               lr=self.lr,
                               weight_decay=0.01 # Default
                               )
        
        self.scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
            optimizer,
            mode='min',
            factor=0.1,
            patience=2,
            min_lr=1e-6,
            verbose=True
        )

        return optimizer
    

        
    #def optimizer_step(self, epoch=None,
    #                   batch_idx=None,
    #                   optimizer=None,
    #                   optimizer_idx=None,
    #                   optimizer_closure=None,
    #                   on_tpu=None,
    #                   using_native_amp=None,
    #                   using_lbfgs=None,
    #                   second_order_closure=None):              
    #    
    #    if batch_idx == 0: # to call the scheduler after each validation
    #        
    #        self.scheduler.step(self.epoch_val_loss)
    #        
    #        print(f'metric: {self.epoch_val_loss}, \
    #              best: {self.scheduler.best}, \
    #                  num_bad_epochs: {self.scheduler.num_bad_epochs}') # for debugging
    #    
    #    optimizer.step()
    #    
    #    optimizer.zero_grad()

Thank you!
hey @tsuijenk
optimizer _closure must be passed in optimizer.step() since it runs training_step and backward call. You can check the docstrings and examples here: https://pytorch-lightning.readthedocs.io/en/latest/common/lightning_module.html#optimizer-step