not really much I can help you with except just noticed,
you do 5 non-strided convolutions in the beginning and residuals get 1x256x256x256, but that’s a lot. For example all basic ResNet models have stride 2 @ 7x7 conv, followed by maxpool 2 exactly for this reason - too much memory consumption. ResNets convert 3x224x224 to 64x56x56 before first residual block.