INNER CODE UNIT · Python
compute_gradients
horovod/horovod · horovod/_keras/__init__.py:103
def compute_gradients(self, loss, var_list, tape=None):
return self._compute_gradients(loss, var_list, None, tape)
def _compute_gradients(self, loss, var_list, grad_loss=None, tape=None):
"""
Compute gradients of all trainable variables.
See Optimizer.get_gradients() for more info.
In DistributedOptimizer, get_gradients() is overriden to also
allreduce the gradients before returning them.
"""
base_class = super(self.__class__, self)
if _PRE_TF_2_4_0:
return base_class._compute_gradients(
loss, var_list, grad_loss, tape)
tape = tf.GradientTape() if tape is None else tape