Learning rate scheduler not updated in train_step #12
Description
My tensorflow version is 2.1.0. I found that when calling step() of the learning rate scheduler, lr is not updated (the scheduler works fine when tested individually). I guess it has something to do with distributed strategy run process. The problem is fixed if moving the learning rate updating process to the main loop, instead of in training step function.