 // y is fairly small. For large y, it would probably be more efficient
 // to also tile across y.
 template <typename T>
-__global__ void ConvShiftForward(const T *x, const T *y, T *out, int x_width,
-                                 int y_width, int y_half_width,
-                                 int batch_size) {
+__global__ void ConvShiftForward(const T *x, const T *y, int x_width,
+                                 int y_width, int y_half_width, int batch_size,
+                                 T *out) {
   extern __shared__ T mem[];
 
   int tx = threadIdx.x;


    if (dX) {
      T *dx_data = dX->mutable_data<T>(context.GetPlace());
      cudaMemsetAsync(dx_data, 0, dX->numel() * sizeof(T), stream);
      ConvShiftGradX<T><<<grid_dim, x_per_block, 0, stream>>>(
          dout_data, y_data, dx_data, x_width, y_width, y_half_width,
          batch_size);
    }
    if (dY) {
      T *dy_data = dY->mutable_data<T>(context.GetPlace());
      cudaMemsetAsync(dy_data, 0, dY->numel() * sizeof(T), stream);
      ConvShiftDy<T><<<grid_dim, x_per_block, 0, stream>>>(
          x_data, dout_data, dy_data, x_width, y_width, y_half_width,
          batch_size);
    }
  }
};

    int y_width = Y->dims()[1];
    int y_half_width = (y_width - 1) / 2;

    auto &device_ctx = context.cuda_device_context();
    math::SetConstant<platform::GPUPlace, T> zero;

    const int x_per_block = 256;
    int num_x_blocks = DivUp(x_width, x_per_block);
