Blank segmentation map for UNet - Can it be a bug in Pytorch?

[ptrblck]

It seems he already solved this issue: UNet implementation