cosmic-cortex / pytorch-UNet

2D and 3D UNet implementation in PyTorch.
MIT License
147 stars 40 forks source link

fix: activation for last layer of UNet2D #16

Open hardik01shah opened 2 months ago

hardik01shah commented 2 months ago

Softmax was being applied by default to the output of the Last2D block, and the softmax argument in the forward method was unused. Fixed it to use the argument to append a softmax layer only when the argument is True.

Last2D expects wrong input channels, only works when conv depths are multiplied by 2 at each depth. Fixed it to work when this condition does not hold.