I just got Segmentation fault (core dumped)
layer1 input shape : 2, 256, 208, 208
layer2 input shape : 2, 256, 104, 104
layer3 input shape : 2, 256, 52, 52
layer4 input shape : 2, 256, 26, 26
layer5 input shape : 2, 256, 26, 26
And when i tried with batch size 1 instead of 2 it is working fine. the problem starts with batch size > 1 and i also tried reducing the input size to check if does not explode the RAM but still the problem is same.