def METHOD_NAME():
    model = resnet18()
    plugin = TorchFSDPPlugin()
    booster = Booster(plugin=plugin)
    optimizer = SGD(model.parameters(), lr=0.001, momentum=0.9)
    criterion = lambda x: x.mean()
    fsdp_model, optimizer, criterion, _, _ = booster.boost(model, optimizer, criterion)
    inputs = torch.randn(4, 3, 224, 224)
    outputs = None
    def run_model():
        nonlocal outputs
        outputs = fsdp_model(inputs)
        optimizer.zero_grad()
        criterion(outputs).backward()
        optimizer.step()
    with shared_tempdir() as tempdir:
        model_ckpt_path = f"{tempdir}/model"
        optim_ckpt_path = f"{tempdir}/optimizer"
        run_model()
        booster.save_model(fsdp_model, model_ckpt_path, shard=False)
        booster.save_optimizer(optimizer, optim_ckpt_path, shard=False)
        full_msd = fsdp_model.state_dict()
        #full_osd = FSDP.full_optim_state_dict(fsdp_model, optimizer)
        sharded_osd = optimizer.state_dict()
        import copy
        sharded_osd = copy.deepcopy(sharded_osd)
        run_model()
        full_msd_updated = fsdp_model.state_dict()
        #full_osd_updated = FSDP.full_optim_state_dict(fsdp_model, optimizer, rank0_only=True)
        sharded_osd_updated = optimizer.state_dict()
        assert not compare_nested_dict(sharded_osd, sharded_osd_updated)
        assert not compare_nested_dict(full_msd_updated, full_msd)
        outputs_first = fsdp_model(inputs)
        assert criterion(outputs_first) != criterion(outputs)
        booster.load_model(fsdp_model, model_ckpt_path)
        booster.load_optimizer(optimizer, optim_ckpt_path)
        full_msd_restore = fsdp_model.state_dict()
        #full_osd_restore = FSDP.full_optim_state_dict(fsdp_model, optimizer, rank0_only=True)
        sharded_osd_restore = optimizer.state_dict()
        assert compare_nested_dict(sharded_osd, sharded_osd_restore)
        assert compare_nested_dict(full_msd_restore, full_msd)
        outputs_sec = fsdp_model(inputs)
        assert criterion(outputs_sec) == criterion(outputs)