So I looked into this a bit more and found some interesting stuff:

With my 40M parameter model, the memory used is increasing from approx 160MB to approx 640MB, so a factor of 4.
This factor of 4 is the same factor for a model with half the parameters. Just the amount of allocated memory also gets halved, i.e. approx. 85MB to 340MB.
Using SGD instead of ADAM decreases this factor to 2.
Dumping the tensors (using ADAM) before the first epoch returns 40M numbers consisting only of parameters. Dumping the tensors after one epoch results in 120M numbers consisting of parameters and tensors.
Doing the same using SGD only shows 40M numbers consisting of only parameters during all epochs, however memory still increases by a factor of 2 as mentioned above.

This leads me to believe that there is something going with model initialization which I don’t understand. I suspect that the model needs 320MB of memory (because ADAM needs another 320MB to keep track of learning rates for weights and biases). But this doesn’t explain why before training but after initialization the model needs 160MB of memory but after training it needs 320MB. What exactly is being held in memory when I delete everything else? Maybe it’s gradients? But aren’t these save in the optimizer? And wouldn’t gradients also need 320MB?
I do not believe that there is a bug somewhere and that all of this is normal behaviour but I would like to understand what exactly is happening here. Could anyone maybe explain?