How to time operations in a tf.function?

Sorry for the dumb question but I’m very new to TF, I want to time a particular section of code inside a tf.function:

@tf.function
def some_func():
    # other code
    start = time.time()
    # code I want to time
    end = time.time()
    print(end-start)
    # other code

it works for the first couple of function calls but it stops printing anything after that and I’m not sure why.

I tried tf.print() as well but it just gives me the same time again and again so I assume it’s reusing the variable.

There is reliable way with:

But I had a bug tracked at: