How do I get current training step?

Hi,

I’m trying to pass in current training step for my customized scheduler. How should I do this? I noticed that self.global_step will count in validation steps, which is not what I want. I only want to count the training step, so I can do something like momentum_schedule[current_training_step] to get correct scheduling value.