class torch.distributed.autograd.context
Context object to wrap forward and backward passes when using
   distributed autograd. The "context_id" generated in the "with"
   statement  is required to uniquely identify a distributed backward
   pass on all workers. Each worker stores metadata associated with
   this "context_id", which is required to correctly execute a
   distributed autograd pass.
Example::
      >>> import torch.distributed.autograd as dist_autograd
      >>> with dist_autograd.context() as context_id:
      >>>     t1 = torch.rand((3, 3), requires_grad=True)
      >>>     t2 = torch.rand((3, 3), requires_grad=True)
      >>>     loss = rpc.rpc_sync("worker1", torch.add, args=(t1, t2)).sum()
      >>>     dist_autograd.backward(context_id, [loss])
torch.distributed.autograd.get_gradients(context_id: int) -> Dict[Tensor, Tensor]
Retrieves a map from Tensor to the appropriate gradient for that
   Tensor accumulated in the provided context corresponding to the