Writing generalisable code for different schedulers

I’m trying to write a class that I will use as a base for future projects. I would like to know how I can generalise the train loop for different schedulers (ones that need to be called after every batch and ones that are called after every epoch).
Thank you