class_docstr: This class is to extract pytorch model topology graph by tracing
class_docstr_tok: ["This","class","is","to","extract","pytorch","model","topology","graph","by","tracing"]
code: def _trace(self, model, dummy_input):
      training = model.training
      model.eval()
      kw_args = {}
      if torch.__version__ >= "1.6.0":
          # only pytorch with version greater than 1.6.0 has the strict option
          kw_args["strict"] = False
      try:
          import pytorch_lightning as pl
      except ImportError:
          is_lightning_module = False
      else:
          if isinstance(model, pl.LightningModule):
              is_lightning_module = True
          else:
              is_lightning_module = False
      if is_lightning_module:
          self.trace = model.to_torchscript(
              method="trace", example_inputs=dummy_input, **kw_args
          )
      else:
          self.trace = torch.jit.trace(model, dummy_input, **kw_args)
      torch._C._jit_pass_inline(self.trace.graph)
      model.train(training)
  
imports: import logging
  import queue
  import re
  from collections import defaultdict
  from typing import List, Dict
  import torch
  from torch.utils.tensorboard._pytorch_graph import NodePy, NodePyIO, NodePyOP, GraphPy
label_desc: Use image node_id 2 for calling the TorchGraph obj's underlying member method code with example usage: obj._trace(model, dummy_input) without return types