def on_train_start(trainer: Trainer, pl_module: LightningModule) -> None:
        """Log model graph to respective logger at training start.

        This method is called automatically at the start of training. For W&B logger,
        it sets up model watching with graph logging enabled.

        Args:
            trainer (Trainer): PyTorch Lightning trainer instance containing logger references.
            pl_module (LightningModule): Lightning module instance to be logged.

        Example:
            >>> from anomalib.callbacks import GraphLogger
            >>> callback = GraphLogger()
            >>> # Called automatically by trainer
            >>> # callback.on_train_start(trainer, model)
        """
        for logger in trainer.loggers:
            if isinstance(logger, AnomalibWandbLogger):
                # NOTE: log graph gets populated only after one backward pass. This won't work for models which do not
                # require training such as Padim
                logger.watch(pl_module, log_graph=True, log="all")
                break
Log model graph to respective logger at training start.

        This method is called automatically at the start of training. For W&B logger,
        it sets up model watching with graph logging enabled.

        Args:
            trainer (Trainer): PyTorch Lightning trainer instance containing logger references.
            pl_module (LightningModule): Lightning module instance to be logged.

        Example:
            >>> from anomalib.callbacks import GraphLogger
            >>> callback = GraphLogger()
            >>> # Called automatically by trainer
            >>> # callback.on_train_start(trainer, model)
        
src/anomalib/callbacks/graph.py