Hi,
If you need gradients all the way back to the first iteration, it is expected that the memory usage is going to be quite high as it needs to keep track of all the buffers to be able to compute the gradients.

Does Gumbel-softmax of each iteration cause such growing memory usage?

It is all the intermediary results that are required to compute the gradients that are expensive. You can try to use checkpoint 39 to reduce the memory usage at the cost of more compute.

Could I use  torch.cuda.empty_cache()  after each epoch to release memory usage?

That won’t reduce the memory usage as the memory in the cache is reused to allocate more Tensor. The only thing this will do is slow down your code by forcing emptying and refilling the cache a lot.