Ran into a similar error during one of my experiments. After some digging, noticed that I updated the number of epochs in the training for-loop, but forgot to re-instantiate the scheduler with the new NUM_EPOCHS (which was being defined in another notebook cell earlier).
Instantiating the scheduler with the updated number of epochs, fixed the issue for me.
Here’s how I am setting the scheduler and the training loop with the NUM_EPOCHS
NUM_EPOCHS = 5

# set the scheduler with OneCycleLR
scheduler = torch.optim.lr_scheduler.OneCycleLR(optimizer=optim,max_lr=0.01,steps_per_epoch=len(train_loader), epochs=NUM_EPOCHS)

# training loop
for epoch in range(NUM_EPOCHS):
    model.train() # set model in training mode
    ....