       } else {
         auto* logits_data = logits->data<T>();
         auto* labels_data = labels->data<int64_t>();
-        HardLabelSoftmaxWithCrossEntropy<T>(
-            context.cuda_device_context(), logits_data, labels_data, loss_data,
-            softmax_data, n, d, axis_dim, ignore_index);
+        if (ignore_index >= 0 && ignore_index < axis_dim) {
+          SoftmaxWithCrossEntropyHardLabel<T, true>(
+              context.cuda_device_context(), rank, axis, logits_data,
+              labels_data, loss_data, softmax_data, n, axis_dim, d / axis_dim,
+              ignore_index);
+        } else {
+          SoftmaxWithCrossEntropyHardLabel<T, false>(
+              context.cuda_device_context(), rank, axis, logits_data,
+              labels_data, loss_data, softmax_data, n, axis_dim, d / axis_dim,
+              ignore_index);
+        }
       }
     }
   }

      SoftCrossEntropyGradientKernel<T><<<grid, block, 0, stream>>>(
          logit_grad_data, loss_grad_data, label_data, n, d, remain);
    } else {
      int64_t grid = (n * remain + block - 1) / block;
      const int64_t* label_data = labels->data<int64_t>();
      CrossEntropyGrad<T><<<grid, block, 0, stream>>>(
          logit_grad_data, label_data, n, d, remain, ignore_index);
      int64_t num = n * d;
      grid = (num + block - 1) / block;
      Scale<T><<<grid, block, 0, stream>>>(logit_grad_data, loss_grad_data, num,
                                           d, remain, label_data, ignore_index);
    }
  }
};

      } else {  // HardLabel
        auto* logits_data = softmax->data<T>();
        auto* labels_data = labels->data<int64_t>();
        int threads = 128;
        int blocks = (n * d / axis_dim + threads - 1) / threads;
        if (ignore_index >= 0 && ignore_index < axis_dim) {
          CrossEntropyHardLabel<T, true><<<
              blocks, threads, 0, context.cuda_device_context().stream()>>>(
              loss_data, logits_data, labels_data, n, axis_dim, d / axis_dim,
              ignore_index);
        } else {
          CrossEntropyHardLabel<T, false><<<
              blocks, threads, 0, context.cuda_device_context().stream()>>>(
              loss_data, logits_data, labels_data, n, axis_dim, d / axis_dim,
              ignore_index);
        }
      }

      // cause of input is softmax
