     MT g = static_cast<MT>(grad[id]);
     MT mom1 = static_cast<MT>(moment1[id]);
     MT mom2 = static_cast<MT>(moment2[id]);
+
+    p *= (static_cast<MT>(1.0) - lr * coeff);
+
     mom1 = beta1 * mom1 + (static_cast<MT>(1.0) - beta1) * g;
     mom2 = beta2 * mom2 + (static_cast<MT>(1.0) - beta2) * g * g;
-    p -= lr_orig * coeff * p;
-    p -= lr * (mom1 /
-               (sqrt(mom2) + epsilon * sqrt(static_cast<MT>(1.0) - beta2_pow)));
+
+    MT denom = (sqrt(mom2) / sqrt(static_cast<MT>(1.0) - beta2_pow)) + epsilon;
+
+    p += (mom1 / denom) * (-(lr / (static_cast<MT>(1.0) - beta1_pow)));
 
     moment1_out[id] = mom1;
     moment2_out[id] = mom2;

    int ndim) {
  int id = blockIdx.x * blockDim.x + threadIdx.x;
  MT lr = *lr_ * lr_ratio;
  MT lr_orig = lr;

  lr *= sqrt(static_cast<MT>(1.0) - beta2_pow) /
        (static_cast<MT>(1.0) - beta1_pow);

  for (; id < ndim; id += blockDim.x * gridDim.x) {
    auto row_idx =

    if (lazy_mode && row_idx < 0) {
      return;
    } else {
      MT mom1 = static_cast<MT>(mom1_[id]);
      MT mom2 = static_cast<MT>(mom2_[id]);

      MT p = master_param ? master_param[id] : static_cast<MT>(param_[id]);
      MT g = row_idx >= 0
                 ? static_cast<MT>(grad_[row_idx * row_numel + id % row_numel])
                 : static_cast<MT>(0);

      p *= (static_cast<MT>(1.0) - lr * coeff);

      mom1 = beta1 * mom1 + (static_cast<MT>(1.0) - beta1) * g;
      mom2 = beta2 * mom2 + (static_cast<MT>(1.0) - beta2) * g * g;

      MT denom =
          (sqrt(mom2) / sqrt(static_cast<MT>(1.0) - beta2_pow)) + epsilon;

      p += (mom1 / denom) * (-(lr / (static_cast<MT>(1.0) - beta1_pow)));

      // Write back to global memory
      mom1_out_[id] = mom1;
