def setup(self, trainer: pl.Trainer, pl_module: pl.LightningModule, stage: str | None = None) -> None:
        """Initialize NNCF compression when training begins.

        This method is called when training or testing begins. It wraps the PyTorch model
        using the NNCF compression controller to prepare it for compression during training.

        Args:
            trainer (pl.Trainer): PyTorch Lightning trainer instance
            pl_module (pl.LightningModule): The Lightning module containing the model to compress
            stage (str | None, optional): Current stage of training. Defaults to ``None``.
        """
        del stage  # `stage` variable is not used.

        if self.nncf_ctrl is not None:
            return

        # Get validate subset to initialize quantization,
        # because train subset does not contain anomalous images.
        init_loader = InitLoader(trainer.datamodule.val_dataloader())
        config = register_default_init_args(self.config, init_loader)

        self.nncf_ctrl, pl_module.model = wrap_nncf_model(
            model=pl_module.model,
            config=config,
            dataloader=trainer.datamodule.train_dataloader(),
            init_state_dict=None,  # type: ignore[arg-type]
        )
Initialize NNCF compression when training begins.

        This method is called when training or testing begins. It wraps the PyTorch model
        using the NNCF compression controller to prepare it for compression during training.

        Args:
            trainer (pl.Trainer): PyTorch Lightning trainer instance
            pl_module (pl.LightningModule): The Lightning module containing the model to compress
            stage (str | None, optional): Current stage of training. Defaults to ``None``.
        
src/anomalib/callbacks/nncf/callback.py