class_docstr: Log model graph to respective logger.
class_docstr_tok: ["Log","model","graph","to","respective","logger","."]
docstr: Unwatch model if configured for wandb and log it model graph in Tensorboard if specified.
  
  Args:
      trainer: Trainer object which contans reference to loggers.
      pl_module: LightningModule object which is logged.
docstr_tok: ["Unwatch","model","if","configured","for","wandb","and","log","it","model","graph","in","Tensorboard","if","specified",".","Args",":","trainer",":","Trainer","object","which","contans","reference","to","loggers",".","pl_module",":","LightningModule","object","which","is","logged","."]
code: def on_train_end(
      self, trainer: Trainer, pl_module: LightningModule
  ) -> None:
      """Unwatch model if configured for wandb and log it model graph in Tensorboard if specified.
  
      Args:
          trainer: Trainer object which contans reference to loggers.
          pl_module: LightningModule object which is logged.
      """
  
      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)
  
imports: import torch
  from pytorch_lightning import Callback, LightningModule, Trainer
  from anomalib.utils.loggers import AnomalibCometLogger, AnomalibTensorBoardLogger, AnomalibWandbLogger
label_desc: Use image node_id 2 for calling the GraphLogger obj's underlying member method code with example usage: obj.on_train_end(trainer, pl_module) without return types