rvandeghen:

After 2-3 epochs it stays around the same value whatever the train-eval-inference phase.


In that case I might have misunderstood the issue and thought the memory would increase in each eproch.
If you are seeing some memory increase in the first epochs, this might be due to memory fragmentation. E.g. if you are finishing the first training and validation run, PyTorch might free some intermediate tensors, which are not referenced anymore. The next memory allocation might not fit in the freed blocks, so that new memory has to be allocated.
How large is the memory increase after the first, second and third epoch?