You’re absolutely right! I found that out just before I saw your reply, the network seems to learn properly now. Looks like I missed the line


nn.Parameter  - A kind of Tensor, that is  automatically registered as a parameter when assigned as an attribute to a   Module .

when reading the Blitz Tutorial 1. Endet up costing me about two days of work to figure that out…



 ptrblck:

Also, instead of sigmoid + nn.BCELoss , I would recommend to use raw logits + nn.BCEWithLogitsLoss for more numerical stability.


I have switched to a prediction of size 2 (because I want to do binary segmentation) along Dim1 + nn.CrossEntropyLoss, but I will try your suggestion and see if it can further improve the results! I guess sticking with BCELoss would save some resources in the decoder.
The only thing that is not quite clear to me is why the model seems to only have ~790k trainable parameters instead of the ~2M it had in Keras, despite using the same number of filters in each layer (obtained using sum(p.numel() for p in net.parameters() if p.requires_grad) )