code: def plot_model(
      model: ClassyModel,
      size: Tuple[int, ...] = (3, 224, 224),
      input_key: Optional[Union[str, List[str]]] = None,
      writer: Optional["SummaryWriter"] = None,
      folder: str = "",
      train: bool = False,
  ) -> None:
      """Visualizes a model in TensorBoard.
  
      The TensorBoard writer can be either specified directly via `writer` or can
      be specified via a `folder`.
  
      The model can be run in training or evaluation model via the `train` argument.
  
      Example usage on devserver:
       - Install TensorBoard using: `sudo feature install tensorboard`
       - Start TensorBoard using: `tensorboard --port=8098 --logdir <folder>`
      """
  
      assert (
          writer is not None or folder != ""
      ), "must specify SummaryWriter or folder to create SummaryWriter in"
      input = get_model_dummy_input(model, size, input_key)
      if writer is None:
          writer = SummaryWriter(log_dir=folder, comment="Model graph")
      with torch.no_grad():
          orig_train = model.training
          model.train(train)
          writer.add_graph(model, input_to_model=(input,))
          model.train(orig_train)
          writer.flush()
  
file_path: ClassyVision/classy_vision/generic/visualize.py
imports: import math
  from typing import Any, Callable, Dict, List, Optional, Tuple, Union
  import numpy
  import torch
  import torch.nn.modules
  from classy_vision.generic.util import get_model_dummy_input, is_pos_int
  from classy_vision.models import ClassyModel
  from PIL import Image
label_desc: Use image node_id 5 for calling a global function with example usage: plot_model(model, size, input_key, writer, folder, train) without return types