Batch Size of 5 takes 7GB on my UNet model, Can anyone help me free up some space on the GPU

I’ve tried running it but it’s hitting assertion in line assert out.size(2) == prev_x.size(2). What’s the shape of the input that you are using?

input size (572x572) exactly as it is the paper
I pushed my entire code on git It’s been running all night on our lab’s devbox but I only got to 100 epoch on training size of 3000 images :confused:

Thank you so much

at batch size of 7, it’s taking 5GB for me.

 5675MiB / 16273MiB

You wrote the network pretty efficiently, i dont think there’s memory improvements possible.

Try adding this flag for speed improvments:

1 Like