scheduler.step() should be right after optimizer.step() and both should be in the training loop.
for epoch in range(N_EPOCHS):
   for batch_i, (X, y) in train_dataloader:
     optimizer.step() 
     scheduler.step()