   int tid = threadIdx.x + blockIdx.x * blockDim.x;
   const cooperative_groups::grid_group cg = cooperative_groups::this_grid();
   for (int i = 0; i < op_num; ++i) {
-    int numel = lars_warpper->numel_arr[i];
+    int numel = lars_warpper.numel_arr[i];
     MT param_norm = static_cast<MT>(0);
     MT grad_norm = static_cast<MT>(0);
-    L2NormKernel<T, MT>(&cg, lars_warpper->p_arr[i], lars_warpper->g_arr[i],
-                        p_buffer, g_buffer, numel, lars_warpper->repeat_arr[i],
+    L2NormKernel<T, MT>(&cg, lars_warpper.p_out_arr[i], lars_warpper.g_arr[i],
+                        p_buffer, g_buffer, numel, lars_warpper.repeat_arr[i],
                         rescale_grad, 0, &param_norm, &grad_norm);
     MomentumUpdate<T, MT>(
-        lars_warpper->p_arr[i], lars_warpper->g_arr[i],
-        lars_warpper->v_out_arr[i], lars_warpper->p_out_arr[i],
-        lars_warpper->v_out_arr[i], lars_warpper->master_p_arr[i],
-        lars_warpper->master_p_out_arr[i], lars_warpper->lr_arr[i], mu,
-        lars_warpper->weight_decay_arr[i], lars_coeff, epsilon, rescale_grad,
+        lars_warpper.p_out_arr[i], lars_warpper.g_arr[i],
+        lars_warpper.v_out_arr[i], lars_warpper.p_out_arr[i],
+        lars_warpper.v_out_arr[i], lars_warpper.master_p_out_arr[i],
+        lars_warpper.master_p_out_arr[i], lars_warpper.lr_arr[i], mu,
+        lars_warpper.weight_decay_arr[i], lars_coeff, epsilon, rescale_grad,
         param_norm, grad_norm, tid, grid_stride, numel, is_amp);
   }
 }

};

template <typename T, typename MT>
__global__ void MergedMomentumLarsKernel(LarsParamWarpper<T, MT>* lars_warpper,
                                         MT* __restrict__ p_buffer,
                                         MT* __restrict__ g_buffer,
                                         const int op_num, const MT mu,

        size_t temp_numel = param[i]->numel();
        total_numel += temp_numel;
        lars_warpper.numel_arr[i] = temp_numel;
        lars_warpper.g_arr[i] = grad[i]->data<T>();
        lars_warpper.lr_arr[i] = learning_rate[i]->data<MT>();
        lars_warpper.p_out_arr[i] =
            param_out[i]->mutable_data<T>(ctx.GetPlace());
        lars_warpper.v_out_arr[i] =
            velocity_out[i]->mutable_data<MT>(ctx.GetPlace());
        lars_warpper.weight_decay_arr[i] = static_cast<MT>(weight_decay_arr[i]);
        PADDLE_ENFORCE_EQ(
            param[i]->data<T>(), lars_warpper.p_out_arr[i],
            platform::errors::InvalidArgument(
                "Input(Param) and Output(ParamOut) must be the same Tensors."));
        PADDLE_ENFORCE_EQ(velocity[i]->data<MT>(), lars_warpper.v_out_arr[i],
                          platform::errors::InvalidArgument(
                              "Input(Velocity) and Output(VelocityOut) must be "
                              "the same Tensors."));
      }
      int64_t avg_numel = total_numel / op_num;
      LarsThreadConfig<float> lars_thread_config(avg_numel, sm_num,
