diff --git a/easy_rec/python/loss/zero_inflated_lognormal.py b/easy_rec/python/loss/zero_inflated_lognormal.py
index 8124edc0c..da1e03d25 100644
--- a/easy_rec/python/loss/zero_inflated_lognormal.py
+++ b/easy_rec/python/loss/zero_inflated_lognormal.py
@@ -49,6 +49,8 @@ def zero_inflated_lognormal_loss(labels, logits, name=''):
   """
   loss_name = name if name else 'ziln_loss'
   labels = tf.cast(labels, dtype=tf.float32)
+  if labels.shape.ndims == 1:
+    labels = tf.expand_dims(labels, 1)  # [B, 1]
   positive = tf.cast(labels > 0, tf.float32)
 
   logits = tf.convert_to_tensor(logits, dtype=tf.float32)