


        optim.step()
  >>>     # All ranks reach here without hanging/erroring




static notify_join_context(joinable)
  Notifies the join context manager that the calling process has
  not yet joined; then, if "throw_on_early_termination=True",
  checks if uneven inputs have been detected (i.e. if one process
  has already joined) and throws an exception if so.

  This method should be called from a "Joinable" object before its
  per-iteration collective communications. For example, this
  should be called at the beginning of the forward pass in
  "DistributedDataParallel".

  Only the first "Joinable" object passed into the context manager
  performs the collective communications in this method, and for
  the others, this method is vacuous.

  Parameters:
     **joinable** (*Joinable*) -- the "Joinable" object calling
     this method.

  Returns:
     An async work handle for the all-reduce meant to notify the
