Thank you very much for the links. I was thinking I could be a bit more imaginative and could use the thresholding in the loss. However, the thresholding turns into a step function which gives infinite gradient at 0.5 and 0 elsewhere. That is why I was getting zero or none gradient.
This is my updated version with non-zero gradient during back propagation.
class Dice_Loss(torch.nn.Module):
def init(self):
super().init()
def forward(self, output, target, smooth=1):
pred_mask = nn.functional.sigmoid(output)
true_pos = torch.sum(torch.mul(pred_mask, target))
true_pos = 2*true_pos
    # just in case there is not 1 in the prediction or target
    # stop overfitting as well
    true_pos = true_pos + smooth
    
    pred_square = torch.sum(torch.mul(pred_mask, pred_mask))
    target_square = torch.sum(torch.mul(target, target)) 
    return 1-(true_pos/(pred_square+target_square + smooth))