How to Track GPU Memory Allocation and Release for Each Tensor During DNN Training in PyTorch?

Hello everyone,

I’m currently working on GPU memory management and I would like to know how to track the memory allocated and released by each tensor throughout the DNN training process in PyTorch. From what I understand, torch.profiler can only provide memory usage at the kernel level, not at the tensor level. So, how can I achieve this?

Are you looking for something like this Understanding GPU Memory 1: Visualizing All Allocations over Time | PyTorch

Thank you very much!!!