ptrblck:

torch.cuda.memory_summary()


Thanks for the help ! I deleted the python cache of my project  (I forgot to delete it when I changed gpu) and the problem seems to have gone more or less, I only have a 1GiB difference during training after the first two epochs compared to with the two rtx2080ti. Using torch.cuda.memory_summary(), we can still observe some memory usage higher than what the rtx2080ti could handle (a peak memory usage of 22Ggib for example), but I expect this is just the optimizer enjoying itself with the extra memory.