       if (soft_label) {
         auto* logits_data = softmax->data<T>();
         auto* labels_data = labels->data<T>();
-        CrossEntropyFusedKernel(logits_data, labels_data, loss_data, n, d,
-                                axis_dim,
-                                context.cuda_device_context().stream());
+
+        const int kDimLog2 = static_cast<int>(Log2Ceil(axis_dim));
+        const int kDimCeil = 1 << kDimLog2;
+#ifdef __HIPCC__
+        int kThreadPerBlock = 256;
+#else
+        int kThreadPerBlock = 512;
+#endif
+        int kBatchPerBlock = 1;
+        int blocks = (n * d + kBatchPerBlock - 1) / kBatchPerBlock;
+        dim3 threads(kThreadPerBlock / kBatchPerBlock, kBatchPerBlock, 1);
+
+        CrossEntropySoftLabel<T, T, false><<<
+            blocks, threads, 0, context.cuda_device_context().stream()>>>(
+            loss_data, NULL, logits_data, labels_data, n, axis_dim,
+            d / axis_dim, kDimLog2);
       } else {  // HardLabel
         auto* logits_data = softmax->data<T>();
         auto* labels_data = labels->data<int64_t>();

  int blocks = (batch_size + batches_per_block - 1) / batches_per_block;
  dim3 threads(kWarpSize, warps_per_block, 1);

  switch (Log2Elements) {
    SOFTMAX_WARP_FORWARD_CASE(0, T, AccT);
    SOFTMAX_WARP_FORWARD_CASE(1, T, AccT);
    SOFTMAX_WARP_FORWARD_CASE(2, T, AccT);

  using AccT = typename details::MPTypeTrait<T>::Type;

  // use 128 threads per block to maximimize gpu utilization
  const int log2_elements = static_cast<int>(Log2Ceil(element_count));
  const int kDimCeil = 1 << log2_elements;
  int kWarpSize = (kDimCeil < 32) ? kDimCeil : 32;
  int batches_per_warp = (kDimCeil <= 128) ? 2 : 1;
  constexpr int threads_per_block = 128;
