Can't run backpropagation through my function

Mehdi_Ibrahimli: <p>Hi everyone I am new to PyTorch and I have a simple training loop:</p>
<pre><code class="lang-auto">for epoch in range(num_epochs):                   
    for i,data in enumerate(training_set):       
        optimizer.zero_grad()                             
        output_layer = PIAE_model(data)
        loss = criterion(output_layer,data)
        loss.backward()                            
        optimizer.step()                          
        loss_log.append(loss.item())
</code></pre>
<p>and the models methods look like this</p>
<pre><code class="lang-auto">def inverter(self,x):
        x = self.up(self.act(self.bn1(self.conv1(x))))     # convolution, batch normalization... 
        x = self.up(self.act(self.bn2(self.conv2(x))))    
        x = self.up((self.bn3(torch.abs(self.conv3(x)))))   
        return x

   def findif(self,vel_model):
       for n in range(0,self.nt):
                d2 = []
                sourceINJ = []
                
                for i in range(self.nx):  # for all the grid points                    
                    d2.append((self.p[i + 1].detach() - 2 * self.p[i].detach() + self.p[i - 1].detach()) / self.dx ** 2 ) 
                    sourceINJ.append(self.p[i].detach())

                sourceINJ[1] = sourceINJ[1] + self.q[n]
                
                self.d2pdx2 = torch.stack(d2)
                self.p = torch.stack(sourceINJ)

            
           self.p = 2 * self.p.detach() + vel_model ** 2 * self.dt * self.d2pdx2.detach()
           self.traces[n] = self.p[1]
        return self.traces
    
    def forward(self, x):
        vel_model = self.inverter(x)
        seis_model = self.findif(vel_model)
        return seis_model
</code></pre>
<p>In the findif method I am using list append to avoid in place operations d2 list calculates the second spatial derivative and sourceINJ does <code>self.p[1] = self.p[1]+self.q[n]</code>. The variables <code>self.p</code> is a torch 0 tensor at first and after the first iteration becomes a functions of <code>vel_model</code> and <code>self.d2pdx2</code> is a function of <code>self.p</code>. And <code>self.dt</code> is a scalar 1.</p>
<p>So the problem here lies in this line <code>self.p = 2 * self.p.detach() + vel_model ** 2 * self.dt * self.d2pdx2.detach()</code> in this form gradients are not computed properly and I get NaNs in my loss and setting torch.autograd.set_detect_anomaly(True) raises ‘Function ‘CudnnBatchNormBackward0’ returned nan values in its 0th output.’</p>
<p>replacing the plus sign in the formula with the multiplication <code>2 * self.p.detach() * vel_model ** 2 * self.dt * self.d2pdx2.detach()</code> lets the training run smoothly. removing the <code>self.d2pdx2</code> and the coefficient 2 such that <code>self.p.detach() + vel_model ** 2 * self.dt</code> also runs smoothly.</p>
<p>Please help me circumvent this problem.</p>
<p>I tried using list append as before</p>
<pre><code class="lang-auto">test = []
for i in range(self.nx):
    test.append(2 * self.p[i].detach() + vel_model[i] ** 2 * self.dt * self.d2pdx2[i].detach())
self.p = torch.stack(test)
</code></pre>
<p>But this didn’t help, I also tried cloning some tensors but no success.</p>
<p>Thank you for reading</p>