dict is only populated on rank 0, and all other ranks return
        an empty "dict".
  Warning:

    Unlike "torch.optim.Optimizer.state_dict()", this method uses
    full parameter names as keys instead of parameter IDs.

  Note:

    Like in "torch.optim.Optimizer.state_dict()", the tensors
    contained in the optimizer state dict are not cloned, so there
    may be aliasing surprises. For best practices, consider saving
    the returned optimizer state dict immediately, e.g. using
    "torch.save()".

  Parameters:
     * **model** (*torch.nn.Module*) -- Root module (which may or
       may not be a "FullyShardedDataParallel" instance) whose
       parameters were passed into the optimizer "optim".

     * **optim** (*torch.optim.Optimizer*) -- Optimizer for
       "model" 's parameters.

     * **optim_input**
       (*Optional**[**Union**[**List**[**Dict**[**str**,
