tf.tape.gradient() returns None for dice loss

Question

I am trying to implement semantic image segmentation with num of classes = 21. The model(U-net) predicts a mask with dimensions (batch_size, H, W, num of classes). I want to use the Dice loss function. When I try to use tf.tape.gradient() to fetch the gradients, it is returning None values. Can you help me to debug the code?

@tf.autograph.experimental.do_not_convert
def dice_loss(pred_mask, gt_mask, smooth=1e-5):
    
    #num_classes = pred_mask.shape[-1]
    pred_mask = tf.keras.layers.Flatten()(pred_mask)
    gt_mask = tf.keras.layers.Flatten()(gt_mask)
    intersection = tf.reduce_sum(gt_mask * pred_mask)
    mask_sum = tf.reduce_sum(gt_mask) + tf.reduce_sum(pred_mask)
    dice = 2 * (intersection + smooth)/(mask_sum + smooth)
    dice = tf.reduce_mean(dice)
    
    return 1 - dice
    

def ohe_pred_mask(pred_mask):    
    
    num_classes = pred_mask.shape[-1]
    ohe_mask = tf.argmax(pred_mask, axis = -1)
    ohe_mask = tf.one_hot(ohe_mask, depth = num_classes)
    
    return ohe_mask

@tf.autograph.experimental.do_not_convert
def train(model, learning_rate, epochs):
    
    DL = dataloader(project_folder = root,
                    image_size = (256, 256),
                    mode = 'train')
    
    ds = DL.tf_dataset(batch = 10)
    
    if model == 'Unet':
        model = Unet(name = "U_Net")
    
    optimizer = tf.compat.v1.train.AdamOptimizer(learning_rate = learning_rate)
    
    for epoch in range(epochs):
        for images, masks in ds.take(1):
            with tf.GradientTape() as tape:
                pred = model(images)
                pred = ohe_pred_mask(pred)
                loss = dice_loss(pred, masks)
                
            grads = tape.gradient(loss, model.trainable_variables)
            print(grads)