code: def __init__(self, module):
      super(DistributedDataParallel, self).__init__()
      self.warn_on_half = (
          True if dist._backend == dist.dist_backend.GLOO else False
      )
  
      self.module = module
      param_list = [
          param
          for param in self.module.state_dict().values()
          if torch.is_tensor(param)
      ]
      if dist._backend == dist.dist_backend.NCCL:
          for param in param_list:
              assert (
                  param.is_cuda
              ), "NCCL backend only supports model parameters to be on GPU."
  
      # broadcast parameters
      flat_dist_call(param_list, dist.broadcast, (0,))
  
      # all reduce gradient hook
      def allreduce_params():
          if self.needs_reduction:
              self.needs_reduction = False
          else:
              return
          grads = [
              param.grad.data
              for param in self.module.parameters()
              if param.grad is not None
          ]
          flat_dist_call(grads, dist.all_reduce)
  
      for param in list(self.module.parameters()):
  
          def allreduce_hook(*unused):
              param._execution_engine.queue_callback(allreduce_params)
  
          if param.requires_grad:
              param.register_hook(allreduce_hook)
  
file_path: DALI/docs/examples/use_cases/video_superres/common/distributed.py
imports: import torch
  from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors
  import torch.distributed
  from torch.nn.modules import Module
label_desc: Use image node_id 1 to create a new DistributedDataParallel object from inherited base classes: Module with example: obj = DistributedDataParallel(module)