Just in case if anyone will have the same error after you load state dicts to re-train/continue to train model, you need to update total_steps value in scheduler state dict before you load it:
scheduler = OneCycleLR(optimizer,...steps_per_epoch=int(len(train_loader)),epochs=args.epochs)
sch_dict = checkpoint['scheduler']
sch_dict['total_steps'] = sch_dict['total_steps'] + args.epochs * int(len(train_loader))
scheduler.load_state_dict(sch_dict)

I think it’s related to only OneCycleLR.