Thanks for the code!
I’ve formatted it for you. You can add code snippets using three backticks `.
I’m not sure, if torch.cuda.memory_allocated() synchronizes all CUDA calls, so could you run your code again with torch.cuda.synchronize() before each memory check?
This will make sure that all CUDA calls (which are asynchronous by default) are done before checking the memory usage.