def METHOD_NAME(self):
    orig_model = CompositeParamModel(device=torch.device("cuda"))
    composable_model = copy.deepcopy(orig_model)
    fully_shard(composable_model, policy=ModuleWrapPolicy({UnitModule}))
    composable_optim = _optim_cls(composable_model.parameters(), lr=_optim_lr)
    orig_model = FSDP(orig_model)
    orig_optim = _optim_cls(orig_model.parameters(), lr=_optim_lr)
    self._test_optim_state_save_load(
        orig_model, orig_optim, composable_model, composable_optim
    )