This helped, at least now I know I have a memory issue, getting a not enough memory error. I wouldn’t have expected that a network taking less than 8GB of GPU memory for forward prop would need more than 60 for backwards. A 3D volume of size 200,200,4 instead of 256,256,4 seems to work on both CPU and GPU.
This is also puzzling me, as running on CPU it takes up to 70% of my ram to perform it, so at least 30GB, while the GPU has absolutely no issue with this computation.
What is going on, and is there a way around this other than rescaling the input?