Performance issue in train.py

[DLPerf]

Hello! Our static bug checker has found a performance issue in train.py: Python type argument ep is passed to tf.function decorated function lr_sch instead of tensors. As there is a for loop, every time the value changes, a new graph will be generated, and that can trigger tf.function retracing warning.

Here is the tensorflow document to support it.

Looking forward to your reply.

[DLPerf]

Maybe you can just pass ep as tf.constant(ep).
We are investigating this kind of issues, and your answer will be of great help to our work. Thank you in advance! @Katexiang