def METHOD_NAME(self, model: Union[Model, Lazy[Model]]) -> Model:
    if isinstance(model, Lazy):
        model = model.construct()
    if not self.fsdp_config.move_params_to_cpu:
        model.to(self.train_config.worker_local_default_device)
    return FSDP(model, **self.fsdp_config.as_kwargs())