Found the culprit thanks to the memory stats. In a torch::autograd::Function, when I create a tensor during forward and store it in ctx->saved_data (for example ctx->saved_data["some_tensor"] = some_tensor; in forward, where some_tensor contains intermediary results I need to reuse later in the backward method), when we arrive at the backward method the tensor doesn’t get dereferenced even after backward finishes. However, if I add a ctx->saved_data["some_tensor"] = torch::Tensor(); at the end of backward (explicitly dereferencing some_tensor), I don’t have any memory leaks anymore.
@ptrblck can you please confirm whether or not it is expected behavior that tensors inside ctx do not get dereferenced even after the end of backward ?