Hi thank you for your response. Are you suggesting to use CrossEntropyLoss? My model and training steps are same as in the following link: UNET model on GPU runtime error
The issue is for each image patch that I have a very small portion of it belongs to mask pixel class of 1. My ground truth only has 2-pixel values. Either 1 or zero.
I am not sure how to implement a weighted loss increasing it for the loss class of 0. Can you suggest some methods?