Unfortunately, it is pretty difficult to do it outside of my model because there are many computations performed with the data before actually handing them to torch.cholesky

What about just logging the size of the matrix and try with a matrix with random values?

When you say “Are you sure that you don’t hold on to some memory across iterations?” what do you mean by that? What kind of functions or coding could lead to holding onto memory?

This can happen mostly if store data in a data structure at each iteration, always appending. Accumulate across iteration (like a epoch_loss) Tensors that still require gradients and hold onto their history