Performance issue about tf.function

[DLPerf]

Hello! Our static bug checker has found a performance issue in adahessian_tf/run_experiments.py and adahessian_tf/cifar_training_tools.py: cifar_training is repeatedly called in a for loop, but there is a tf.function decorated function step defined and called in cifar_training.

In that case, when cifar_training is called in a loop, the function step will create a new graph every time, and that can trigger tf.function retracing warning.

Here is the tensorflow document to support it.

Briefly, for better efficiency, it's better to use:

@tf.function
def inner():
    pass

def outer():
    inner()  

than:

def outer():
    @tf.function
    def inner():
        pass
    inner()

Looking forward to your reply.

[DLPerf]

Since there is model here, it is not easy to make a change, do you have any idea? @amirgholami @nestordemeure @sIncerass @yaozhewei