def on_train_end(trainer: Trainer, pl_module: LightningModule) -> None:
        """Log model graph at training end and cleanup.

        This method is called automatically at the end of training. It:
        - Logs the model graph for TensorBoard and Comet loggers
        - Unwatches the model for W&B logger

        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_end(trainer, model)
        """
        for logger in trainer.loggers:
            if isinstance(logger, AnomalibCometLogger | AnomalibTensorBoardLogger):
                logger.log_graph(pl_module, input_array=torch.ones((1, 3, 256, 256)))
            elif isinstance(logger, AnomalibWandbLogger):
                logger.experiment.unwatch(pl_module)
Log model graph at training end and cleanup.

        This method is called automatically at the end of training. It:
        - Logs the model graph for TensorBoard and Comet loggers
        - Unwatches the model for W&B logger

        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_end(trainer, model)
        
src/anomalib/callbacks/graph.py