         invvar,
         n1, n2,
         grad_input);
+    } else {
+      cuComputeGradInput<T, U, false, true, false><<<blocks1, threads1, nshared, stream>>>(
+        dout,
+        input,
+        output,
+        gamma,
+        beta,
+        mean,
+        invvar,
+        n1, n2,
+        grad_input);
+    }
+  } else {
+    if (gamma == nullptr) {
+      cuComputeGradInput<T, U, true, false, simplified><<<blocks1, threads1, nshared, stream>>>(
+        dout,
+        input,
+        output,
+        gamma,
+        beta,
+        mean,
+        invvar,
+        n1, n2,
+        grad_input);
+    } else {
+      cuComputeGradInput<T, U, true, true, simplified><<<blocks1, threads1, nshared, stream>>>(
+        dout,
+        input,
+        output,
+        gamma,
+        beta,
+        mean,
+        invvar,
+        n1, n2,
+        grad_input);
+    }
   }
 }
 

  }
}

template <typename T, typename U, bool use_mean, bool simplified>
__global__ void cuComputeGradInput(
    const T* __restrict__ dout,
    const T* __restrict__ input,

  int nshared =
      threads1.y > 1 ? threads1.y * threads1.x * sizeof(U) : 0;
  if (mean == nullptr && !simplified) {
    if (gamma == nullptr) {
      cuComputeGradInput<T, U, false, false, false><<<blocks1, threads1, nshared, stream>>>(
        dout,
        input,
        output,
