code: import torch
  from torch.nn.modules.pooling import MaxPool2d
  
  from .activation import ReLU6, Hardswish, ELU, LeakyReLU, Sigmoid
  from .batchnorm import BatchNorm2d, BatchNorm3d
  from .normalization import LayerNorm, GroupNorm, InstanceNorm1d, \
      InstanceNorm2d, InstanceNorm3d
  from .conv import _ConvNd, Conv1d, Conv2d, Conv3d
  from .conv import ConvTranspose1d, ConvTranspose2d, ConvTranspose3d
  from .linear import Linear
  from .embedding_ops import Embedding, EmbeddingBag
  
  from .functional_modules import FloatFunctional, FXFloatFunctional, QFunctional
  
  
  class Quantize(torch.nn.Module):
      r"""Quantizes an incoming tensor
  
      Args:
       `scale`: scale of the output Quantized Tensor
       `zero_point`: zero_point of output Quantized Tensor
       `dtype`: data type of output Quantized Tensor
  
      Attributes:
        `scale`, `zero_point`, `dtype`
  
      Examples::
          >>> t = torch.tensor([[1., -1.], [1., -1.]])
          >>> scale, zero_point, dtype = 1.0, 2, torch.qint8
          >>> qm = Quantize(scale, zero_point, dtype)
          >>> qt = qm(t)
          >>> print(qt)
          tensor([[ 1., -1.],
                  [ 1., -1.]], size=(2, 2), dtype=torch.qint8, scale=1.0, zero_point=2)
      """
  
      scale: torch.Tensor
      zero_point: torch.Tensor
  
      def __init__(self, scale, zero_point, dtype):
          super(Quantize, self).__init__()
          self.register_buffer('scale', torch.tensor([scale]))
          self.register_buffer('zero_point', torch.tensor([zero_point], dtype=torch.long))
          self.dtype = dtype
  
      def forward(self, X):
          return torch.quantize_per_tensor(X, float(self.scale),
                                           int(self.zero_point), self.dtype)
  
      @staticmethod
      def from_float(mod):
          assert hasattr(mod, 'activation_post_process')
          scale, zero_point = mod.activation_post_process.calculate_qparams()
          return Quantize(scale.float().item(), zero_point.long().item(), mod.activation_post_process.dtype)
  
      def extra_repr(self):
          return 'scale={}, zero_point={}, dtype={}'.format(self.scale, self.zero_point, self.dtype)
  
  
  class DeQuantize(torch.nn.Module):
      r"""Dequantizes an incoming tensor
  
      Examples::
          >>> input = torch.tensor([[1., -1.], [1., -1.]])
          >>> scale, zero_point, dtype = 1.0, 2, torch.qint8
          >>> qm = Quantize(scale, zero_point, dtype)
          >>> quantized_input = qm(input)
          >>> dqm = DeQuantize()
          >>> dequantized = dqm(quantized_input)
          >>> print(dequantized)
          tensor([[ 1., -1.],
                  [ 1., -1.]], dtype=torch.float32)
      """
  
      def __init__(self):
          super(DeQuantize, self).__init__()
  
      def forward(self, Xq):
          return Xq.dequantize()
  
      @staticmethod
      def from_float(mod):
          return DeQuantize()
  
  __all__ = [
      'BatchNorm2d',
      'BatchNorm3d',
      '_ConvNd',
      'Conv1d',
      'Conv2d',
      'Conv3d',
      'ConvTranspose1d',
      'ConvTranspose2d',
      'ConvTranspose3d',
      'DeQuantize',
      'ELU',
      'Embedding',
      'EmbeddingBag',
      'GroupNorm',
      'Hardswish',
      'InstanceNorm1d',
      'InstanceNorm2d',
      'InstanceNorm3d',
      'LayerNorm',
      'LeakyReLU',
      'Linear',
      'MaxPool2d',
      'Quantize',
      'ReLU6',
      'Sigmoid',
      # Wrapper modules
      'FloatFunctional',
      'FXFloatFunctional',
      'QFunctional',
  ]
   import itertools
  from typing_extensions import Protocol
  import warnings
  
  import torch
  from ..parameter import is_lazy
  
  
  class _LazyProtocol(Protocol):
      """This is to avoid errors with mypy checks for
      The attributes in a mixin:
      https://mypy.readthedocs.io/en/latest/more_types.html#mixin-classes
      """
      def _register_load_state_dict_pre_hook(self, hook):
          ...
  
      def register_forward_pre_hook(self, hook):
          ...
  
      def _lazy_load_hook(
              self, state_dict, prefix, local_metadata, strict,
              missing_keys, unexpected_keys, error_msgs):
          ...
  
      def _get_name(self):
          ...
  
      def _infer_parameters(self, module, input):
          ...
  
      @property
      def _parameters(self):
          ...
  
      @property
      def _buffers(self):
          ...
  
      @property
      def _non_persistent_buffers_set(self):
          ...
  
      @property
      def _load_hook(self):
          ...
  
      @property
      def _initialize_hook(self):
          ...
  
  
  class LazyModuleMixin:
      r"""A mixin for modules that lazily initialize parameters, also known as "lazy modules."
  
      .. warning:
          Lazy modules are an experimental new feature under active development,
          and their API is likely to change.
  
      Modules that lazily initialize parameters, or "lazy modules",
      derive the shapes of their parameters from the first input(s)
      to their forward method. Until that first forward they contain
      :class:`torch.nn.UninitializedParameter` s that should not be accessed
      or used, and afterward they contain regular :class:`torch.nn.Parameter` s.
      Lazy modules are convenient since they don't require computing some
      module arguments, like the :attr:`in_features` argument of a
      typical :class:`torch.nn.Linear`.
  
      After construction, networks with lazy modules should first
      be converted to the desired dtype and placed on the expected device.
      This is because lazy modules only perform shape inference so the usual dtype
      and device placement behavior applies.
      The lazy modules should then perform "dry runs" to initialize all the components in the module.
      These "dry runs" send inputs of the correct size, dtype, and device through
      the network and to each one of its lazy modules. After this the network can be used as usual.
  
      >>> class LazyMLP(torch.nn.Module):
      ...    def __init__(self):
      ...        super().__init__()
      ...        self.fc1 = torch.nn.LazyLinear(10)
      ...        self.relu1 = torch.nn.ReLU()
      ...        self.fc2 = torch.nn.LazyLinear(1)
      ...        self.relu2 = torch.nn.ReLU()
      ...
      ...    def forward(self, input):
      ...        x = self.relu1(self.fc1(input))
      ...        y = self.relu2(self.fc2(x))
      ...        return y
      >>> # constructs a network with lazy modules
      >>> lazy_mlp = LazyMLP()
      >>> # transforms the network's device and dtype
      >>> # NOTE: these transforms can and should be applied after construction and before any 'dry runs'
      >>> lazy_mlp = mlp.cuda().double()
      >>> lazy_mlp
      LazyMLP( (fc1): LazyLinear(in_features=0, out_features=10, bias=True)
        (relu1): ReLU()
        (fc2): LazyLinear(in_features=0, out_features=1, bias=True)
        (relu2): ReLU()
      )
      >>> # performs a dry run to initialize the network's lazy modules
      >>> lazy_mlp(torch.ones(10,10).cuda())
      >>> # after initialization, LazyLinear modules become regular Linear modules
      >>> lazy_mlp
      LazyMLP(
        (fc1): Linear(in_features=10, out_features=10, bias=True)
        (relu1): ReLU()
        (fc2): Linear(in_features=10, out_features=1, bias=True)
        (relu2): ReLU()
      )
      >>> # attaches an optimizer, since parameters can now be used as usual
      >>> optim = torch.optim.SGD(mlp.parameters(), lr=0.01)
  
      A final caveat when using lazy modules is that the order of initialization of a network's
      parameters may change, since the lazy modules are always initialized after other modules.
      For example, if the LazyMLP class defined above had a :class:`torch.nn.LazyLinear` module
      first and then a regular :class:`torch.nn.Linear` second, the second module would be
      initialized on construction and the first module would be initialized during the first dry run.
      This can cause the parameters of a network using lazy modules to be initialized differently
      than the parameters of a network without lazy modules as the order of parameter initializations,
      which often depends on a stateful random number generator, is different.
      Check :doc:`/notes/randomness` for more details.
  
      Lazy modules can be serialized with a state dict like other modules. For example:
  
      >>> lazy_mlp = LazyMLP()
      >>> # The state dict shows the uninitialized parameters
      >>> lazy_mlp.state_dict()
      OrderedDict([('fc1.weight', Uninitialized parameter),
                   ('fc1.bias',
                    tensor([-1.8832e+25,  4.5636e-41, -1.8832e+25,  4.5636e-41, -6.1598e-30,
                             4.5637e-41, -1.8788e+22,  4.5636e-41, -2.0042e-31,  4.5637e-41])),
                   ('fc2.weight', Uninitialized parameter),
                   ('fc2.bias', tensor([0.0019]))])
  
  
      Lazy modules can load regular :class:`torch.nn.Parameter` s (i.e. you can serialize/deserialize
      initialized LazyModules and they will remain initialized)
  
  
      >>> full_mlp = LazyMLP()
      >>> # Dry run to initialize another module
      >>> full_mlp.forward(torch.ones(10, 1))
      >>> # Load an initialized state into a lazy module
      >>> lazy_mlp.load_state_dict(full_mlp.state_dict())
      >>> # The state dict now holds valid values
      >>> lazy_mlp.state_dict()
      OrderedDict([('fc1.weight',
                    tensor([[-0.3837],
                            [ 0.0907],
                            [ 0.6708],
                            [-0.5223],
                            [-0.9028],
                            [ 0.2851],
                            [-0.4537],
                            [ 0.6813],
                            [ 0.5766],
                            [-0.8678]])),
                   ('fc1.bias',
                    tensor([-1.8832e+25,  4.5636e-41, -1.8832e+25,  4.5636e-41, -6.1598e-30,
                             4.5637e-41, -1.8788e+22,  4.5636e-41, -2.0042e-31,  4.5637e-41])),
                   ('fc2.weight',
                    tensor([[ 0.1320,  0.2938,  0.0679,  0.2793,  0.1088, -0.1795, -0.2301,  0.2807,
                              0.2479,  0.1091]])),
                   ('fc2.bias', tensor([0.0019]))])
  
      Note, however, that the loaded parameters will not be replaced when doing a "dry run" if they are initialized
      when the state is loaded. This prevents using initialized modules in different contexts.
      """
  
      # modules inheriting from this will change their __class__ to the specified
      # one after they are fully initialized
      cls_to_become = None
  
      def __init__(self: _LazyProtocol, *args, **kwargs):
          # Mypy doesnt like this super call in a mixin
          super().__init__(*args, **kwargs)  # type: ignore
          self._load_hook = self._register_load_state_dict_pre_hook(self._lazy_load_hook)
          self._initialize_hook = self.register_forward_pre_hook(self._infer_parameters)
          warnings.warn('Lazy modules are a new feature under heavy development '
                        'so changes to the API or functionality can happen at any moment.')
  
      def _save_to_state_dict(self: _LazyProtocol, destination, prefix, keep_vars):
          # This should be ideally implemented as a hook,
          # but we should override `detach` in the UninitializedParameter to return itself
          # which is not clean
          for name, param in self._parameters.items():
              if param is not None:
                  if not (is_lazy(param) or keep_vars):
                      param = param.detach()
                  destination[prefix + name] = param
          for name, buf in self._buffers.items():
              if buf is not None and name not in self._non_persistent_buffers_set:
                  if not (is_lazy(buf) or keep_vars):
                      buf = buf.detach()
                  destination[prefix + name] = buf
  
      def _lazy_load_hook(
              self: _LazyProtocol, state_dict, prefix, local_metadata, strict,
              missing_keys, unexpected_keys, error_msgs):
          """load_state_dict pre-hook function for lazy buffers and parameters.
  
          The purpose of this hook is to adjust the current state and/or
          ``state_dict`` being loaded so that a module instance serialized in
          both un/initialized state can be deserialized onto both un/initialized
          module instance.
          See comment in ``torch.nn.Module._register_load_state_dict_pre_hook``
          for the details of the hook specification.
          """
          for name, param in itertools.chain(self._parameters.items(), self._buffers.items()):
              key = prefix + name
              if key in state_dict and param is not None:
                  input_param = state_dict[key]
                  if is_lazy(param):
                      # The current parameter is not initialized but the one being loaded one is
                      # create a new parameter based on the uninitialized one
                      if not is_lazy(input_param):
                          with torch.no_grad():
                              param.materialize(input_param.shape)
  
      def initialize_parameters(self: _LazyProtocol, *args, **kwargs):
          r"""Initialize parameters according to the input batch properties.
          This adds an interface to isolate parameter initialization from the
          forward pass when doing parameter shape inference.
          """
          raise NotImplementedError('initialize_parameters is not implemented for {}'.format(self.__class__.__name__))
  
      def has_uninitialized_params(self: _LazyProtocol):
          r"""Check if a module has parameters that are not initialized
          """
          # This is to avoid the JIT to track this parameter and force
          # custom modules __setstate__ to add it
          params = self._parameters.values()
          buffers = self._buffers.values()
          for param in itertools.chain(params, buffers):
              if is_lazy(param):
                  return True
          return False
  
      def _infer_parameters(self: _LazyProtocol, module, input):
          r"""Infers the size and initializes the parameters according to the
          provided input batch.
          Given a module that contains parameters that were declared inferrable
          using :class:`torch.nn.parameter.ParameterMode.Infer`, runs a forward pass
          in the complete module using the provided input to initialize all the parameters
          as needed.
          The module is set into evaluation mode before running the forward pass in order
          to avoid saving statistics or calculating gradients
          """
          module.initialize_parameters(*input)
          if module.has_uninitialized_params():
              raise RuntimeError('module {} has not been fully initialized'.format(self._get_name()))
          module._initialize_hook.remove()
          module._load_hook.remove()
          delattr(module, '_initialize_hook')
          delattr(module, '_load_hook')
          if module.cls_to_become is not None:
              module.__class__ = module.cls_to_become
  
  
      def _replicate_for_data_parallel(self: _LazyProtocol):
          raise RuntimeError('Modules with uninitialized parameters can\'t be used with `DataParallel`. '
                             'Run a dummy forward pass to correctly initialize the modules')
   from collections import OrderedDict
  from typing import Dict, Any
  
  from tensorboard.compat.proto.config_pb2 import RunMetadata
  from tensorboard.compat.proto.graph_pb2 import GraphDef
  from tensorboard.compat.proto.step_stats_pb2 import StepStats, DeviceStepStats
  from tensorboard.compat.proto.versions_pb2 import VersionDef
  
  import torch
  from ._proto_graph import node_proto
  
  methods_OP = ['attributeNames', 'hasMultipleOutputs', 'hasUses', 'inputs',
                'kind', 'outputs', 'outputsSize', 'scopeName']
  # Some additional methods to explure for methods_IO are
  #
  #   'unique' (type int)
  #   'type' (type <Tensor<class 'torch._C.Type'>>)
  #
  # But the below are sufficient for now.
  methods_IO = ['node', 'offset', 'debugName']
  
  GETATTR_KIND = 'prim::GetAttr'
  CLASSTYPE_KIND = 'ClassType'
  
  class NodeBase(object):
      def __init__(self, debugName=None, inputs=None, scope=None, tensor_size=None, op_type='UnSpecified', attributes=''):
          # TODO; Specify a __slots__ for this class or potentially
          # used namedtuple instead
          self.debugName = debugName
          self.inputs = inputs
          self.tensor_size = tensor_size
          self.kind = op_type
          self.attributes = attributes
          self.scope = scope
  
      def __repr__(self):
          repr = []
          repr.append(str(type(self)))
          for m in dir(self):
              if '__' not in m:
                  repr.append(m + ': ' + str(getattr(self, m)) + str(type(getattr(self, m))))
          return '\n'.join(repr) + '\n\n'
  
  
  class NodePy(NodeBase):
      def __init__(self, node_cpp, valid_methods):
          super(NodePy, self).__init__(node_cpp)
          valid_methods = valid_methods[:]
          self.inputs = []
  
          for m in valid_methods:
              if m == 'inputs' or m == 'outputs':
                  list_of_node = list(getattr(node_cpp, m)())
                  io_unique_names = []
                  io_tensor_sizes = []
                  for n in list_of_node:
                      io_unique_names.append(n.debugName())
                      if n.isCompleteTensor():
                          io_tensor_sizes.append(n.type().sizes())
                      else:
                          io_tensor_sizes.append(None)
  
                  setattr(self, m, io_unique_names)
                  setattr(self, m + 'tensor_size', io_tensor_sizes)
  
              else:
                  setattr(self, m, getattr(node_cpp, m)())
  
  
  class NodePyIO(NodePy):
      def __init__(self, node_cpp, input_or_output=None):
          super(NodePyIO, self).__init__(node_cpp, methods_IO)
          try:
              tensor_size = node_cpp.type().sizes()
          except RuntimeError:
              tensor_size = [1, ]  # fail when constant model is used.
          self.tensor_size = tensor_size
          # Kind attribute string is purely descriptive and will be shown
          # in detailed information for the node in TensorBoard's graph plugin.
          #
          # NodePyOP nodes get this from their kind() method.
          self.kind = 'Parameter'
          if input_or_output:
              self.input_or_output = input_or_output
              self.kind = 'IO Node'
  
  
  class NodePyOP(NodePy):
      def __init__(self, node_cpp):
          super(NodePyOP, self).__init__(node_cpp, methods_OP)
          # Replace single quote which causes strange behavior in TensorBoard
          # TODO: See if we can remove this in the future
          self.attributes = str({k: node_cpp[k] for k in node_cpp.attributeNames()}).replace("'", ' ')
          self.kind = node_cpp.kind()
  
  
  class GraphPy(object):
      """Helper class to convert torch.nn.Module to GraphDef proto and visualization
      with TensorBoard.
  
      GraphDef generation operates in two passes:
  
      In the first pass, all nodes are read and saved to two lists.
      One list is for input/output nodes (nodes_io), which only have inbound
      or outbound connections, but not both. Another list is for internal
      operator nodes (nodes_op). The first pass also saves all scope name
      appeared in the nodes in scope_name_appeared list for later processing.
  
      In the second pass, scope names are fully applied to all nodes.
      debugNameToScopedName is a mapping from a node's ID to its fully qualified
      scope name. e.g. Net1/Linear[0]/1. Unfortunately torch.jit doesn't have
      totally correct scope output, so this is nontrivial. The function
      populate_namespace_from_OP_to_IO and find_common_root are used to
      assign scope name to a node based on the connection between nodes
      in a heuristic kind of way. Bookkeeping is done with shallowest_scope_name
      and scope_name_appeared.
      """
      def __init__(self):
          self.nodes_op = []
          self.nodes_io = OrderedDict()
          self.unique_name_to_scoped_name = {}
          self.shallowest_scope_name = 'default'
          self.scope_name_appeared = []
  
      def append(self, x):
          if isinstance(x, NodePyIO):
              self.nodes_io[x.debugName] = x
          if isinstance(x, NodePyOP):
              self.nodes_op.append(x)
  
      def printall(self):
          print('all nodes')
          for node in self.nodes_op:
              print(node)
          for key in self.nodes_io:
              print(self.nodes_io[key])
  
      def find_common_root(self):
          for fullscope in self.scope_name_appeared:
              if fullscope:
                  self.shallowest_scope_name = fullscope.split('/')[0]
  
      def populate_namespace_from_OP_to_IO(self):
          for node in self.nodes_op:
              for node_output, outputSize in zip(node.outputs, node.outputstensor_size):
                  self.scope_name_appeared.append(node.scopeName)
                  self.nodes_io[node_output] = NodeBase(node_output,
                                                        node.inputs,
                                                        node.scopeName,
                                                        outputSize,
                                                        op_type=node.kind,
                                                        attributes=node.attributes)
  
          self.find_common_root()
  
          for node in self.nodes_op:
              for input_node_id in node.inputs:
                  self.unique_name_to_scoped_name[input_node_id] = node.scopeName + '/' + input_node_id
  
          for key, node in self.nodes_io.items():
              if type(node) == NodeBase:
                  self.unique_name_to_scoped_name[key] = node.scope + '/' + node.debugName
              if hasattr(node, 'input_or_output'):
                  self.unique_name_to_scoped_name[key] = node.input_or_output + '/' + node.debugName
  
              if hasattr(node, 'scope') and node.scope is not None:
                  self.unique_name_to_scoped_name[key] = node.scope + '/' + node.debugName
                  if node.scope == '' and self.shallowest_scope_name:
                      self.unique_name_to_scoped_name[node.debugName] = self.shallowest_scope_name + '/' + node.debugName
  
          # replace name
          for key, node in self.nodes_io.items():
              self.nodes_io[key].inputs = [self.unique_name_to_scoped_name[node_input_id] for node_input_id in node.inputs]
              if node.debugName in self.unique_name_to_scoped_name:
                  self.nodes_io[key].debugName = self.unique_name_to_scoped_name[node.debugName]
  
      def to_proto(self):
          """
          Converts graph representation of GraphPy object to TensorBoard
          required format.
          """
          # TODO: compute correct memory usage and CPU time once
          # PyTorch supports it
          nodes = []
          for v in self.nodes_io.values():
              nodes.append(node_proto(v.debugName,
                                      input=v.inputs,
                                      outputsize=v.tensor_size,
                                      op=v.kind,
                                      attributes=v.attributes))
          return nodes
  
  
  def parse(graph, trace, args=None, omit_useless_nodes=True):
      """This method parses an optimized PyTorch model graph and produces
      a list of nodes and node stats for eventual conversion to TensorBoard
      protobuf format.
  
      Args:
        graph (PyTorch module): The model graph to be parsed.
        trace (PyTorch JIT TracedModule): The model trace to be parsed.
        args (tuple): input tensor[s] for the model.
        omit_useless_nodes (boolean): Whether to remove nodes from the graph.
      """
      n_inputs = len(args)
  
      scope = {}
      nodes_py = GraphPy()
      for node in graph.inputs():
          if omit_useless_nodes:
              if len(node.uses()) == 0:  # number of user of the node (= number of outputs/ fanout)
                  continue
  
          if node.type().kind() != CLASSTYPE_KIND:
              nodes_py.append(NodePyIO(node, 'input'))
  
      attr_to_scope: Dict[Any, str] = dict()
      for node in graph.nodes():
          if node.kind() == GETATTR_KIND:
              attr_name = node.s('name')
              parent = node.input().node()
              if parent.kind() == GETATTR_KIND:  # If the parent node is not the top-level "self" node
                  parent_attr_name = parent.s('name')
                  parent_scope = attr_to_scope[parent_attr_name]
                  attr_scope = parent_scope.split('/')[-1]
                  attr_to_scope[attr_name] = '{}/{}.{}'.format(parent_scope, attr_scope, attr_name)
              else:
                  attr_to_scope[attr_name] = '__module.{}'.format(attr_name)
              # We don't need classtype nodes; scope will provide this information
              if node.output().type().kind() != CLASSTYPE_KIND:
                  node_py = NodePyOP(node)
                  node_py.scopeName = attr_to_scope[attr_name]  # type: ignore
                  nodes_py.append(node_py)
          else:
              nodes_py.append(NodePyOP(node))
  
      for i, node in enumerate(graph.outputs()):  # Create sink nodes for output ops
          node_pyio = NodePyIO(node, 'output')
          node_pyio.debugName = "output.{}".format(i + 1)
          node_pyio.inputs = [node.debugName()]
          nodes_py.append(node_pyio)
  
      def parse_traced_name(module):
          if isinstance(module, torch.jit.TracedModule):
              module_name = module._name
          else:
              module_name = getattr(module, 'original_name', "Module")
          return module_name
  
      alias_to_name = dict()
      base_name = parse_traced_name(trace)
      for name, module in trace.named_modules(prefix='__module'):
          mod_name = parse_traced_name(module)
          attr_name = name.split('.')[-1]
          alias_to_name[name] = '{}[{}]'.format(mod_name, attr_name)
  
      for node in nodes_py.nodes_op:
          module_aliases = node.scopeName.split('/')
          replacements = [
              alias_to_name[alias]
              if alias in alias_to_name
              else alias.split('.')[-1]
              for alias in module_aliases
          ]
          node.scopeName = base_name
          if any(replacements):
              node.scopeName += '/' + '/'.join(replacements)
  
      nodes_py.populate_namespace_from_OP_to_IO()
      return nodes_py.to_proto()
  
  
  def graph(model, args, verbose=False):
      """
      This method processes a PyTorch model and produces a `GraphDef` proto
      that can be logged to TensorBoard.
  
      Args:
        model (PyTorch module): The model to be parsed.
        args (tuple): input tensor[s] for the model.
        verbose (bool): Whether to print out verbose information while
          processing.
      """
      with torch.onnx.select_model_mode_for_export(model, torch.onnx.TrainingMode.EVAL):  # TODO: move outside of torch.onnx?
          try:
              trace = torch.jit.trace(model, args)
              graph = trace.graph
              torch._C._jit_pass_inline(graph)
          except RuntimeError as e:
              print(e)
              print('Error occurs, No graph saved')
              raise e
  
      if verbose:
          print(graph)
      list_of_nodes = parse(graph, trace, args)
      # We are hardcoding that this was run on CPU even though it might have actually
      # run on GPU. Note this is what is shown in TensorBoard and has no bearing
      # on actual execution.
      # TODO: See if we can extract GPU vs CPU information from the PyTorch model
      # and pass it correctly to TensorBoard.
      #
      # Definition of StepStats and DeviceStepStats can be found at
      # https://github.com/tensorflow/tensorboard/blob/master/tensorboard/plugins/graph/tf_graph_common/test/graph-test.ts
      # and
      # https://github.com/tensorflow/tensorboard/blob/master/tensorboard/compat/proto/step_stats.proto
      stepstats = RunMetadata(step_stats=StepStats(dev_stats=[DeviceStepStats(device="/device:CPU:0")]))
      return GraphDef(node=list_of_nodes, versions=VersionDef(producer=22)), stepstats
      # The producer version has been reverse engineered from standard
      # TensorBoard logged data.
  
apis: torch.tensor
apis: torch.no_grad
apis: torch.onnx.select_model_mode_for_export torch._C._jit_pass_inline torch.jit.trace