add graph to Tensorboard logger

[Borda]

🚀 Feature

adding a graph to TB logger as it is shown here
https://www.learnopencv.com/tensorboard-with-pytorch-lightning/
maybe also add histograms...

Motivation

[Fan-Luo]

Hi,

I am trying to add the computational graph to tensorboard with
self.logger.experiment.add_graph(TriviaQA(self.args), sampleInput)

but got the error message

File "./.local/lib/python3.6/site-packages/torch/jit/init.py", line 742, in _check_trace
raise TracingCheckError(*diag_info)
torch.jit.TracingCheckError: Tracing failed sanity checks!
ERROR: Graphs differed across invocations!

(~30000 lines after Graph diff)

First diverging operator:
Node diff:
- %9 : torch.torch.nn.modules.linear.___torch_mangle_226.Linear = prim::GetAttrname="linear_sp_sent"
? ^^^
+ %9 : torch.torch.nn.modules.linear.___torch_mangle_470.Linear = prim::GetAttrname="linear_sp_sent"
? ^^^

I tried one solution in the similar issue (pytorch/pytorch#23993) by not declare sampleInput as a variable, but does not work.
The other solution mentioned is to set check_trace to False. But I do not call torch.jit.trace directly.

Another issue (#3003) mentioned one requirement: user has defined the model.example_input_array attribute, but I assume it is only for automatically logging the computational graph, right?

Any other solution I can try?

Thank you

[JakobHavtorn]

It would be nice to be able to pipe arguments for torch.jit.trace through .add_graph() as kwargs such that we can disable the trace checking for stochastic models.