code: def manual_gradient_all_reduce(
      model: DistributedDataParallel,
  ) -> None:
      """
      Gradient reduction function used after backward is done.
      """
      w = []
      for p in model.parameters():
          if p.grad is not None:
              work = dist.all_reduce(
                  p.grad.data, group=model.process_group, async_op=True
              )
              w.append(work)
      for work in w:
          work.wait()
      for p in model.parameters():
          if p.grad is not None:
              p.grad.data.div_(dist.get_world_size())
  
file_path: vissl/vissl/utils/activation_checkpointing.py
imports: from typing import Dict, List
  import torch.distributed
  from torch import nn
  from torch.nn import Module
  from torch.nn.parallel import DistributedDataParallel
label_desc: Use image node_id 3 for calling a global function with example usage: manual_gradient_all_reduce(model) without return types