Calling __exit_ after GradientTape(persistent=True)

Question

I start recording the gradients with tf.GradientTape(persistent=True) in order be able to call gradient_tape.gradient() more than one time. Since this means that the gradients are not being released after calling the gradient() method, do I have to call something like __exit__ at the end of train_step ?

Answer 1

According to https://www.tensorflow.org/api_docs/python/tf/GradientTape:

By default, the resources held by a GradientTape are released as soon as GradientTape.gradient() method is called. To compute multiple gradients over the same computation, create a persistent gradient tape. This allows multiple calls to the gradient() method as resources are released when the tape object is garbage collected.

After train step, as long as no other object has a reference to that tape, the garbage collector will collect it.