StepLR(_LRScheduler) based on batch_sizes

I am going through the transfer learning tutorial, and looked at the source code of StepLR.
It seems if I wanted to change the StepLR to be based on batches instead of epochs,
I just need to run scheduler.step() every x of amount batches, where x is the amount desired and set step_size within lr_scheduler to 1?

That’s actually pretty straightforward, but I can understand not thinking about it, because it’s not explicit from the documentation.

What you want to do is define your StepLR normally in the beginning, but instead of having values correspond to epochs, you set them up w.r.t. batches. Then you call the .step() method of the scheduler at each batch instead of each epoch…

Okay. Yeah I had written that in the OP, but just wanted to make sure. Thank you.