Hi everyone,
I am currently working on a setup where the output of my network is first modified by my own module before computing the loss and backpropagating this, as follows:
class NeuralNet(nn.Module):
    def __init__(self):
        super(NeuralNet, self).__init__()

            self.Layer1 = Layer()
            self.Layer2 = Layer()

            self.Mod = Own_module()

            self.LossLayer = Loss()

    def forward(self, x):
        y1 = self.Layer1(x)
        y2 = self.Layer2(y1)

        self.output = self.Mod(y2)

    def backward(self):
        loss =  self.LossLayer(self.output)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

This works, however, Own_mudule() is an extensive simulation with lots of multiplications and loops (operations are completely fixed). Because of all the computational graphs that need to be saved during this simulation, memory issues occur. I tried solving this using:
self.Mod.requires_grad_(False)

However, this still results in memory issues. And when I use detach() the backpropagation gets blocked completely.
Is there a way I can not save the computational graphs during the Own_module() and essentially bypass it during backpropagation?
Thanks a lot!