 template <typename T>
 __global__ void LabelSmoothGradRunKernel(const int N, const float epsilon,
                                          const T* src, T* dst) {
-  int idx = blockDim.x * blockIdx.x + threadIdx.x;
-  for (; idx < N; idx += blockDim.x * gridDim.x) {
+  CUDA_KERNEL_LOOP(idx, N) {
     dst[idx] = static_cast<T>(1 - epsilon) * src[idx];
   }
 }

__global__ void LabelSmoothRunOriginKernel(const int N, const float epsilon,
                                           const int label_dim, const T* src,
                                           T* dst) {
  int idx = blockDim.x * blockIdx.x + threadIdx.x;
  for (; idx < N; idx += blockDim.x * gridDim.x) {
    dst[idx] = static_cast<T>(1 - epsilon) * src[idx] +
               static_cast<T>(epsilon / label_dim);
  }

__global__ void LabelSmoothRunDistKernel(const int N, const float epsilon,
                                         const int dist_numel, const T* src,
                                         const T* dist_data, T* dst) {
  CUDA_KERNEL_LOOP(idx, N) {
    int dist_idx = idx % dist_numel;
    dst[idx] = static_cast<T>(1 - epsilon) * src[idx] +
               static_cast<T>(epsilon) * dist_data[dist_idx];
