Hi, I was able to solve my problem by changing my steps_per_epoch in the OneCycleLR to 2*(len(train_loader)).
The default configuration of OneCycleLR is set for taking only one step in the training stage per epoch. But since I was training twice per epoch hence it stopped after going through half the times.
I am not sure if you are training multiple times per epoch for adversarial training or something else. You can check into that. Also you need to check how many times scheduler.step() is executed per epoch. If it is excuting more than once per epoch you need to change your step_per_epoch.