Hi,
Maybe you have a bit too many channels? Convs like Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)) on a 512x512x512 image is going to be quite memory heavy !
It will have ((512+2*1-1)/1+1)*((512+2*1-1)/1+1) = 264196 patches each of size 512 x 3x3 = 4608.
Which means that the unfolding for the convolution will contain batch x 1 217 415 168 elements.