def METHOD_NAME(
    self,
    cast_root_forward_inputs_submodule: bool,
    cast_forward_inputs_submodule: bool,
    use_root_no_params: bool,
):
    forward_inputs: Dict[nn.Module, torch.Tensor] = {}
    float16 = MixedPrecision(
        param_dtype=torch.float16,
        cast_root_forward_inputs=cast_root_forward_inputs_submodule,
        cast_forward_inputs=cast_forward_inputs_submodule,
    )
    model = SaveForwardInputsModel(
        forward_inputs=forward_inputs,
        cast_forward_inputs=False,
    ).cuda()
    x = torch.zeros(2, 100, device="cuda")
    # float16 on at least one submodule
    model.c2 = fully_shard(model.c2, mixed_precision=float16)
    # if not `use_root_no_params`, one float32 module will be left for root
    if use_root_no_params:
        # use float16 and wrap all submodules, leaving root without direct parameters
        model.c1 = fully_shard(model.c1, mixed_precision=float16)
        fsdp = fully_shard(model, mixed_precision=float16)
    else:
        fsdp = fully_shard(model)
    return x, fsdp