         out_grad->template mutable_data<T>(),
         &context_);
 
-    for (int i = 0; i < M; ++i) {
-      if (FIRSTDIMS) {
+    if (FIRSTDIMS) {
+      for (int i = 0; i < M; ++i) {
         math::Axpby<T, CUDAContext>(
             N,
             alpha,

  int num_reduce_dims_;
};

template <
    typename T,
    class Context = CUDAContext,

            static_cast<T>(0),
            out_grad->template mutable_data<T>() + i * N,
            &context_);
      }
    } else {
      StripedScaleKernel<T><<<
          CAFFE_GET_BLOCKS(N * M),
          CAFFE_CUDA_NUM_THREADS,
          0,
          context_.cuda_stream()>>>(
          N,
          M,
          grad_in.template data<T>(),
          out_grad->template data<T>(),
          out_grad->template mutable_data<T>());
    }

    return true;
