How to reduce memory usage of backpropagation?

I implemented the [tiny NeRF example](https://github.com/bmild/nerf/blob/master/tiny_nerf.ipynb) 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?

Start with examples/tiny_nerf.rs and compare it with the linked tiny_nerf.ipynb implementation. Use the reported CPU and GPU memory behavior and heaptrack results to trace allocations associated with backpropagation. Done means identifying the implementation cause of the excessive memory use and determining whether the behavior can be reduced or disabled in the relevant parts.

Written by the indexing model from the issue text.