Sorry for late reply.
Yes, I want to go all the way to the first iteration, backprop to i_0 (i.e. input of the network). Additionally, during forward pass, in each iteration, the selection of intermediate feature i_k (i_k can have different size, that means it will not have a constant GPU memory consumption) based on Gumbel-Softmax, which also consumes additional GPU memory (I think).
So, in the forward pass, GPU memory checked by nvidia-smi, will sometimes increase, sometimes be stable, and it can also raise OOM anytime (maybe forward with large intermediate features?)
I want to know, how can I control the GPU memory usage? Does Gumbel-softmax of each iteration cause such growing memory usage? Could I use torch.cuda.empty_cache() after each epoch to release memory usage?
Thanks for you help!