@albanD
Thanks for the feedback.
I just tried to test one iteration of the model with torchviz, and run into an issue.
It seems that it just shows what the graph would be when we compute the model’s output with just one x, (ie. after 1 forward).
In my case I am actually doing 2 forwards of the model, and back-propagating the loss between the 2 outputs of the forward.
Testing it with MNIST I get:
Is there any workaround to test this?
My only concern is that computation might get messed up when the loss is computed between the two forwards (and not against a target).
Thanks!