Sayak_Paul:

But why would you want to compute gradients during inference?


I work in interpretability of machine learning, and most of the methods use the computation of gradients, whether this is done in regards to the inputs or to some feature maps [1, 2].
To calculate the gradients in TensorFlow you need to use the GradientTape and inside this tape you are required to make inference of the model with the  __call__.



 Sayak_Paul:

And this should be utilizing multiple GPUs. I can confirm this from experience.


This is 100% true 
The issue is that you can’t make use of model.predict inside the GradientTape.
That being said, in the second link where there are inference loops with customizations, it really isn’t applicable in this scenario (I think).
Nevertheless, I appreciate all the help and sharing those links. I didn’t know it was possible to do such things in TensorFlow and they will certainly be useful in the future