memory. This API should only be used when the model "state_dict"
      is derived with the context manager "with
      state_dict_type(SHARDED_STATE_DICT):".
  For the detailed usage, refer to "full_optim_state_dict()".

  Warning:

    The returned state dict contains "ShardedTensor" and cannot be
    directly used by the regular "optim.load_state_dict".

  Return type:
     *Dict*[str, *Any*]

static state_dict_type(module, state_dict_type, state_dict_config=None)
  A context manager to set the "state_dict_type" of all the
  descendant FSDP modules of the target module. This context
  manager has the same functions as "set_state_dict_type()". Read
  the document of "set_state_dict_type()" for the detail.

  Example:

     >>> model = DDP(FSDP(...))
     >>> with FSDP.state_dict_type(
     >>>     model,
     >>>     StateDictType.SHARDED_STATE_DICT,
     >>> ):
     >>>     checkpoint = model.state_dict()

  Parameters:
