Hi Mithrandir!
I don’t understand your use case or follow what you are doing in your
code, so I can’t give you a concrete answer.  But I do have some
comments, below.



 mithrandir:

then I apply a data consistency step where I do a conjugate gradient algorithm
…
However, while training the inclusion of the conj grad step is adding hugely to the GPU memory used (without: 2GB, with: 45GB).
This conjugate gradient algorithm […] is an iterative method involving matrix operations.


First, it is perfectly fine to use a loop or iterative algorithm in your
forward pass.  As long as the calculations inside of your loop use
(differentiable) pytorch tensor operations (that support autograd),
autograd and backpropagation will work properly.
As you’ve seen, however, each iteration through the loop will add
another “layer” to the “computation graph,” consuming memory, and
taking time when you backpropagate, perhaps to the extent that you
run out of memory or the computation becomes impractically slow.
Depending on your use case, you may be able to turn off autograd for
your iterative forward function, and write your own corresponding
backward function that requires less memory (and/or time).
The “baseline” use of the conjugate-gradient algorithm is as a linear
solver where the linear operation A . x (with A a matrix, and x a
vector) is given implicitly somehow, or where A is sparse so that
computing A . x is cheaper than computing a full matrix-vector
product or it is impractical to store the full A^-1 in memory.
Let’s say that this is your use case, so that you are solving A . x = b
for x, and you require gradients of the computed x with respect to the
elements of the vector b.
Then x = A^-1 . b, so that the Jacobean (of x with respect to b) is
A^-1.  In backpropagation you typically compute the dot product of
such a Jacobean with a vector of numerical gradients that are being
backpropagated from later stages of your network.  Let’s call this
vector g and the result of the dot product r = A^-1 . g.
r can now be obtained by solving A . r = g for r for which you can
again employ the conjugate-gradient algorithm, taking advantage of
any sparsity in A (or if you need to because the linear operation
A . x is only given implicitly).
You can also write your own backward function, potentially using the
conjugate-gradient algorithm, in the case that you require the gradients
of the computed x with respect to the elements of the matrix A.  Doing
so is relatively straightforward, but may or may not be helpfully efficient
or practical depending on the details of the sparsity structure of A.
Best.
K. Frank