


FSDP.set_state_dict_type(
         >>>     model,
         >>>     StateDictType.SHARDED_STATE_DICT,
         >>>     ShardedStateDictConfig(offload_to_cpu=True),
         >>> )
         >>> checkpoint = model.state_dict()



  Parameters:
     * **module** (*torch.nn.Module*) -- Root module.

     * **state_dict_type** (*StateDictType*) -- the desired
       "state_dict_type" to set.

     * **state_dict_config** (*Optional**[**StateDictConfig**]*)
       -- the configuration for the target "state_dict_type".

  Return type:
     *Tuple*[*StateDictType*, *StateDictConfig*]

static shard_full_optim_state_dict(full_optim_state_dict, model, optim_input=None, optim=None)
  Shards the full optimizer state dict "full_optim_state_dict" by
  remapping the state to flattened parameters instead of
  unflattened parameters and restricting to only this rank's part
  of the optimizer state. The first argument should be the return
