Hi,
To be able to compute the backward pass, the autograd needs to save some intermediary results that are required by the backward formula.
Just given the size of your function, I guess there are a lot of these intermediary states that need to be saved, leading to a significant memory increase.
Note that this is expected and if you use a resnet for example, most of the memory usage will come from the autograd states!
You can try to use checkpoint doc 55 to trade memory for compute: It will discard some of these intermediary results and will recompute them during the backward pass.
So if you wrap a collection of ops inside a checkpoint, you won’t have to save the intermediary states for these. But they will be recomputed during the backward which will be slower.