instance     has been fully initialized, which happens once
         "DistributedDataParallel" gradient buckets have been
         rebuilt.
  Warning:

    This needs to be called on all ranks.

join_hook(**kwargs)
  Returns the ZeRO join hook, which enables training on uneven
  inputs by shadowing the collective communications in the
  optimizer step.

  Gradients must be properly set before this hook is called.

  Parameters:
     **kwargs** (*dict*) -- a "dict" containing any keyword
     arguments to modify the behavior of the join hook at run
     time; all "Joinable" instances sharing the same join context
     manager are forwarded the same value for "kwargs".

  This hook does not support any keyword arguments; i.e. "kwargs"
  is unused.

load_state_dict(state_dict)
  Load the state pertaining to the given rank from the input
  "state_dict", updating the local optimizer as needed.

  Parameters:
