Found the problem. In the resnet I use, there is no pooling layer, just strided convolutions for downsampling while my small custom network uses max pool. By removing max pool and adding strided conv instead, I managed to lower the memory usage to something acceptable. My experiment showed these memory consumptions:

resnet18:

11M parameters
855 MiB of memory with model and batch loaded only
2755 MiB after forward pass
3537 MiB after backward pass


small network with max pool:

400K parameters
811 MiB with model and batch only
4329 MiB after forward
5059 MiB after backward


small network with strided conv:

400K parameters
811 MiB with model and batch only
1615 MiB after forward
1755 MiB after backward



The conclusion to all this is that max pool takes a huge amount of memory. I am not sure it is expected behavior as I don’t know how it affects the computation graph but I will for sure not be using it for my experiments.