[Python] In-place operations in gradient computation in nested PyTorch loops

I'm developing a PyTorch-based model that utilizes nested loops for calculations.

I've encountered a roadblock when attempting to compute gradients. The backpropagation process fails due to what appears to be in-place operations within the loops.

Here's a minimal example illustrating the problem:

import torch

# Setup
length, dx = 1000, 100
num_rows, num_cols = 5, int(length / dx) + 1

# Parameters requiring gradients
C = torch.rand(num_rows, requires_grad=True)
D = torch.rand(num_rows, requires_grad=True)
qobs = torch.rand(num_rows)

# Intermediate calculations
Caux = 1 + C + D
C0, C1, C2 = (-1 + C + D) / Caux, (1 + C - D) / Caux, (1 - C + D) / Caux

# Main computation
qprop = torch.zeros((num_rows, num_cols))
qprop[:, 0] = qobs

for j in range(1, num_cols):
for n in range(1, num_rows):
qprop[n, j] = (C2[n] * qprop[n-1, j-1] +
C1[n] * qprop[n, j-1] +
C0[n] * qprop[n-1, j])

# Attempt gradient computation
loss = qprop.sum()
loss.backward() # This operation fails

Continue reading...