"buffer_dtype" in the first forward pass, and keeps them in
        that dtype thereafter. Model checkpointing always happens in
        full precision.
  * **keep_low_precision_grads** (*bool*) -- This specifies
    whether to upcast gradients back to the full parameter
    precision after the backward pass. This may be set to "False"
    to save memory if using custom optimizers that can perform the
    optimizer step in "reduce_dtype". (Default: "False")

  * **cast_forward_inputs** (*bool*) -- Cast floating point
    tensors in the forward arguments and keyword arguments to
    "param_dtype". (Default: "False")

  * **cast_root_forward_inputs** (*bool*) -- Cast floating point
    tensors in the forward arguments and keyword arguments to
    "param_dtype" for the root FSDP instance. It takes precedence
    over "cast_forward_inputs" for the root FSDP instance.
    (Default: "True")

Note:
 This API is experimental and subject to change.
