Hi Luuk!



 LuukJ:

Is there a way I can not save the computational graphs during the Own_module()


Yes, but …


and essentially bypass it during backpropagation?


You won’t want to bypass your Own_module backpropagation entirely.
Suppose Own_module were simply return -input, that is, it just flipped
the sign of (the gradient of) the loss function?  Then you would be training
your model to maximize the loss function, your predictions would become
increasing poor, and your parameters would likely run off to infinity.
Therefore, you need to backpropagate back through Own_module  with
at least some crude approximation to Own_module's own gradient, at a
minimum getting the sign of its gradient correct at least most of the time.
I would suggest that you package Own_module not as a Module, but
as a torch.autograd.Function.  Your custom Function's forward()
method would implement your “extensive simulation” without saving
the computational graph (using, e.g., .detach() or protected by a
with torch.no_grad(): block), and the computation graph would get
“glued back together” because you would also implement your Function's
.backward() method.  The more accurately .backward() approximates
the actual gradient of .forward(), the better your model will likely train.
As an aside, your NeuralNet Module doesn’t need to (and, to avoid
confusion, shouldn’t) implement a backward() method – in a more
standard design, the optimization code in your .backward() method
would be moved outside of NeuralNet into your training loop after
the application fo NeuralNet to your input data.
Best.
K. Frank