Add PyTorch graph summary to TensorBoard.

pytorch-graph-summary.py

import torch as T

model = T.nn.Sequential()
model.add_module('W0', T.nn.Linear(128, 10))
model.add_module('tanh', T.nn.Tanh())
model.add_module('W1', T.nn.Linear(10, 5))
model.add_module('tanh', T.nn.Tanh())
model.add_module('W2', T.nn.Linear(5, 1))

input = T.randn(16, 128)
    
with T.utils.tensorboard.SummaryWriter('/tmp/log/graph-summary') as summary:
    summary.add_graph(model, x)
    
# Now, run `tensorboard --logdir /tmp/log serve` in shell and open TensorBoard in browser.