code: def average_gradients(model):
      
      size = float(dist.get_world_size())
      for param in model.parameters():
          dist.all_reduce(param.grad.data, op=dist.reduce_op.SUM, group=0)
          param.grad.data /= size
docs: Gradient averaging.