Unfortunately, the documentation for BCEWithLogitsLoss is pretty
much silent about your use case

Wow, thank you for explaining. What you described is honestly something I would not have figured out by myself reading the docs… For this application, using the right loss function and loss weighting was the critical piece. I don’t think the model learning would have converged without it, even in this simplified MVP. My segmentation is now looking great with extremely high accuracy after I implemented the weighting correctly. 

Capture.PNG1913×953 115 KB