Return type:
         Dict[str, Any]
static set_state_dict_type(module, state_dict_type, state_dict_config=None)
  Set the "state_dict_type" and the corresponding (optional)
  configurations of all the descendant FSDP modules of the target
  module. The target module does not have to be a FSDP module. If
  the target module is a FSDP module, its "state_dict_type" will
  also be changed.

  Note:

    This API should be called for only the top-level (root)
    module.

  Note:

    This API enables users to transparently use the conventional
    "state_dict" API to take model checkpoints in cases where the
    root FSDP module is wrapped by another "nn.Module". For
    example, the following will ensure "state_dict" is called on
    all non-FSDP instances, while dispatching into
    *sharded_state_dict* implementation for FSDP:

  Example:

     >>> model = DDP(FSDP(...))
     >>> FSDP.set_state_dict_type(
