def mpi_avg_grads()

in spinup/utils/mpi_pytorch.py [0:0]

• 7 lines of code
• 3 McCabe index (conditional complexity)

def mpi_avg_grads(module):
    """ Average contents of gradient buffers across MPI processes. """
    if num_procs()==1:
        return
    for p in module.parameters():
        p_grad_numpy = p.grad.numpy()   # numpy view of tensor data
        avg_p_grad = mpi_avg(p.grad)
        p_grad_numpy[:] = avg_p_grad[:]