Examples can be found here: https://github.com/pytorch/pytorch/blob/c76fada4a859742ac679013b7428017a782e1432/torch/nn/parallel/distributed.py#L607-L684 7
IIUC, as of today, the communication bucket is still divided by the world size even if the hook is enabled. We are working on removing that division.

  
      github.com
  
  
    pytorch/pytorch/blob/c76fada4a859742ac679013b7428017a782e1432/torch/csrc/distributed/c10d/reducer.cpp#L355-L356 1

auto wrapped =
    c10::scalar_to_tensor(double(1.) / process_group_->getSize());