code: def draw_img_classifier(
      model,
      dataset=None,
      display_param_nodes=False,
      rankdir="TB",
      styles=None,
      input_shape=None,
  ):
      """Draw a PyTorch image classifier to a PNG file.  This a helper function that
      simplifies the interface of draw_model_to_file().
      Args:
          model: PyTorch model instance
          png_fname (string): PNG file name
          dataset (string): one of 'imagenet' or 'cifar10'.  This is required in order to
                            create a dummy input of the correct shape.
          display_param_nodes (boolean): if True, draw the parameter nodes
          rankdir: diagram direction.  'TB'/'BT' is Top-to-Bottom/Bottom-to-Top
                   'LR'/'R/L' is Left-to-Rt/Rt-to-Left
          styles: a dictionary of styles.  Key is module name.  Value is
                  a legal pydot style dictionary.  For example:
                  styles['conv1'] = {'shape': 'oval',
                                     'fillcolor': 'gray',
                                     'style': 'rounded, filled'}
          input_shape (tuple): List of integers representing the input shape.
                               Used only if 'dataset' is None
      """
      dummy_input = distiller.get_dummy_input(
          dataset=dataset,
          device=distiller.model_device(model),
          input_shape=input_shape,
      )
      try:
          non_para_model = distiller.make_non_parallel_copy(model)
          g = SummaryGraph(non_para_model, dummy_input)
  
          return sgraph2dot(g, display_param_nodes, rankdir, styles)
          print("Network PNG image generation completed")
      except FileNotFoundError:
          print(
              "An error has occured while generating the network PNG image."
          )
          print("Please check that you have graphviz installed.")
          print("\t$ sudo apt-get install graphviz")
      finally:
          del non_para_model
  
file_path: tensorwatch/tensorwatch/model_graph/hiddenlayer/pytorch_draw_model.py
imports: import pydot
  from .summary_graph import SummaryGraph
  from .None import distiller
  import torch
  import collections.abc
  import os
label_desc: Use image node_id 2 for calling a global function with example usage: draw_img_classifier(model, dataset, display_param_nodes, rankdir, styles, input_shape) and returns: sgraph2dot