Tresholding the prediction image to binary before sending to loss function

[InnovArul (Arul)]

@einrone

prediction_image[prediction_image >= treshold] = 1.0
prediction_image[prediction_image < treshold] = 0.0

If at all possible, would you be able to try how straight-through estimator works in your case?

The parameters of the model with custom loss function doesn't upgraded thorough its learning over epochs

I am suspecting that torch.where may not be backpropagating gradients. If you really want to threshold, try to use straight-through estimator trick as follows: thresholded_inputs = torch.where(thresholded_inputs < threshold, 0, 1) inputs = (inputs + thresholded_inputs) - inputs.detach() ... calculate IoU loss ...