title: How to Visualize the model graph of a Graph Neural Network in Tensorboard
body: <p>I am trying to visualize the computation graphs of Graph Neural Networks I make to predict properties of Molecules. The model is made in PyTorch and takes as input DGL graphs. The code snippet for trying to visualize the model looks like this:</p>
  <pre class="lang-py prettyprint-override"><code>train_log_dir = f'logs/{datetime.datetime.now().strftime(&quot;%Y%m%d-%H%M%S&quot;)}/train'
  train_summary_writer = tensorboardX.SummaryWriter(train_log_dir)
  train_summary_writer.add_graph(model, [transformer(dataset[0][0]), transformer(dataset[0][0])])
  </code></pre>
  <p>I encounter the following error, TensorBoardX fails to visualize the graph model, refuses to accept DGL graphs as inputs and only wants tensors. Is there any way I can visualize the model?</p>
  <pre><code>RuntimeError: Tracer cannot infer type of (Graph(num_nodes=3, num_edges=4,
        ndata_schemes={'x': Scheme(shape=(10,), dtype=torch.float32)}
        edata_schemes={'w': Scheme(shape=(4,), dtype=torch.float32)}), Graph(num_nodes=3, num_edges=4,
        ndata_schemes={'x': Scheme(shape=(10,), dtype=torch.float32)}
        edata_schemes={'w': Scheme(shape=(4,), dtype=torch.float32)}))
  :Only tensors and (possibly nested) tuples of tensors, lists, or dictsare supported as inputs or outputs of traced functions, but instead got value of type DGLHeteroGraph.
  
  Process finished with exit code 1
  </code></pre>