I was wondering if there is a way to add a custom computation graph to tensorboard using PyTorch. The  website says it works for PyTorch models and tensors. I didn’t find examples where the graphs are being added to tensorboard for custom computation graphs.
So basically my function takes in a couple of tensors, performs calculations within the function creating more tensors and then outputs a result. I wanted to be able to visualize the interactions between the various tensors.
I tried calling writer.add_graph(my_func, (my_args1, my_args2)), but that resulted in an error stating : 'function' object has no attribute 'param_init_net'
I have tried to create an extremely simple working example here.
https://colab.research.google.com/drive/1AIcObBl1QTc3h0hryWo1ZO7ZJ5hk_MfT?usp=sharing 3
Can someone help me out on this?
Thank you.