     AccT log_softmax = func(static_cast<AccT>(logits[tid]));
     softmax[tid] = static_cast<T>(std::exp(log_softmax));
     // loss
-    if (label_valid) {
-      ComputeLoss<T, IgnoreIndex>(loss,
-                                  static_cast<T>(-log_softmax),
-                                  label_id,
-                                  label_value,
-                                  tid,
-                                  1,
-                                  0,
-                                  ignore_index);
-    }
-  }
-
-  // invalid label, write once
-  if (!label_valid && threadIdx.x == 0) {
-    loss[label_id] = static_cast<T>(0.0f);
-  }
-}
-
-template <typename T,
-          typename AccT,
-          typename LabelT,
-          int VecSize,
-          bool IgnoreIndex>
+    ComputeLoss<T>(loss,
+                   static_cast<T>(-log_softmax),
+                   label_id,
+                   label_value,
+                   tid,
+                   1,
+                   0,
+                   ignore_index);
+  }
+}
+
+template <typename T, typename AccT, typename LabelT, int VecSize>
 __global__ void VectorizedSoftmaxForward(T* loss,
                                          T* softmax,
                                          const T* logits,

  int tid = threadIdx.x;
  int label_id = blockIdx.x;
  auto label_value = static_cast<int64_t>(label[label_id]);
  assert(label_value >= 0 && label_value < size || label_value == ignore_index);
  const bool label_valid = label_value >= 0 && label_value < size;
  int loss_id_offset = 0;

  if (offset > 0) {

  // thread ids compute loss[ids] using softmax[idx]
  if (ids < n * d) {
    auto lbl = static_cast<int64_t>(labels[ids]);
    PADDLE_ENFORCE(lbl >= 0 && lbl < dim || lbl == ignore_idx,
                   "The value of label expected >= 0 and < %d, or == %d, "
                   "but got %ld. Please check label value.",
                   dim,
                   ignore_idx,
                   lbl);
    if (lbl == ignore_idx) {
      loss[ids] = static_cast<T>(0.0);
    } else {
      int64_t idx = idx_n * dim * d + lbl * d + idx_d;
      loss[ids] = -Log(softmax[idx]);
    }
  }
}
