def _maybe_reduce(val):
    if (num_distrib() > 1):
        val = val.clone()
        torch.distributed.all_reduce(val, op=torch.distributed.ReduceOp.SUM)
        val /= num_distrib()
    return val
def _maybe_reduce(val):
    if (num_distrib() > 1):
        val = val.clone()
        torch.distributed.all_reduce(val, op=torch.distributed.ReduceOp.SUM)
        val /= num_distrib()
    return val