INNER CODE UNIT · Python
__filtered_reduce_grads
horovod/horovod · horovod/_keras/__init__.py:168
def __filtered_reduce_grads(grads, vars):
rv = []
rg = []
if _IS_TF2:
v2g = {var.ref(): grad for var, grad in zip(vars, grads)}
for var, grad in zip(vars, grads):
if var.ref() not in self._local_vars:
rv.append(var)
rg.append(grad)
else:
v2g = {var: grad for var, grad in zip(vars, grads)}
for var, grad in zip(vars, grads):
if var not in self._local_vars:
rv.append(var)
rg.append(grad)
rg = self._allreduce_grads(rg, rv)
horovod_size = size_op(process_set_id=self.process_set.process_set_id) if int(os.environ.get("HOROVOD_ELASTIC", 0)) else self.process_set.size()