How to reduce memory usage of backpropagation?

[laptou]

I implemented the tiny NeRF example using candle here: https://github.com/laptou/nerfy/blob/fc50dbd61c4012d1f12f556a72474b59a8b3c158/examples/tiny_nerf.rs

The example, which is written using TensorFlow, runs fine on my laptop. My candle implementation consumes all available memory on my laptop, which crashes my desktop session if I use CPU and errors out with a CUDA memory allocation error if I use the GPU. I'm running on a laptop with 32 GB of RAM, 32 GB of swap, and an RTX A3000 w/ 12 GB of VRAM.

I'm barely able to run it on CPU if I decrease the hidden layer size from 256 to 64.

I tracked the memory allocations using heaptrack, and it seems like most of them are related to keeping track of the operations for backpropagation.

Can you spot any obvious issues in my implementation that are causing it to consume so much memory? Is there a way that I can disable or reduce this behavior in some parts of the code to reduce the amount of memory that it uses?