 
   // dJ/dW
   for (int z = 0; z < 4; ++z) {
-    CHECK_CUBLAS_CALL(cublasSetStream(handle, stream[z * 4 + 1]));
+    CHECK_CUBLAS_CALL(cublasSetStream(handle, stream[4 + z]));
     CHECK_CUBLAS_CALL(gemm_gpu<T>(
         handle, CUBLAS_OP_T, CUBLAS_OP_N, K, 5 * D, H * W * N,
         scale, I, K,

      <<<DIV_UP(H * W * N, BLOCK_SIZE1D), BLOCK_SIZE1D>>>(H * W * N, vOnes, 1);
  CHECK_LAST_CUDA_CALL();
  for (int z = 0; z < 4; ++z) {
    CHECK_CUBLAS_CALL(cublasSetStream(handle, stream[z * 4 + 0]));
    CHECK_CUBLAS_CALL(gemv_gpu<T>(
        handle, CUBLAS_OP_T, H * W * N, 5 * D,
        scale, dQ_ptr(z, 0, 0, 0, 0, 0), 6 * D,

        1.0, dW_ptr(z, 0, 0, 0), 5 * D));
  }

  // translate the output tensor in the x-dimension
  T* Oxp = Q + H * W * N;
  kernel_copy_Oxp_to_Q
      <<<DIV_UP(4 * H * W * N * D, BLOCK_SIZE1D), BLOCK_SIZE1D, 0, stream[8]>>>
      (H, W, N, D, O, Oxp);
  // translate the output tensor in the y-dimension
  T* Oyp = Q + H * W * N + 4 * H * W * N * D;
  kernel_copy_Oyp_to_Q
      <<<DIV_UP(4 * H * W * N * D, BLOCK_SIZE1D), BLOCK_SIZE1D, 0, stream[9]>>>
      (H, W, N, D, O, Oyp);
  // wait for data copies
  CHECK_CUDA_CALL(cudaStreamSynchronize(stream[8]));
  CHECK_CUDA_CALL(cudaStreamSynchronize(stream[9]));

  // dJ/dRx
  for (int z = 0; z < 4; ++z) {
    CHECK_CUBLAS_CALL(cublasSetStream(handle, stream[8 + z]));
    CHECK_CUBLAS_CALL(gemm_gpu<T>(
        handle, CUBLAS_OP_T, CUBLAS_OP_N, D, 5 * D, H * W * N,
        scale, Oxp + z * H * W * N * D, D,
        dQ_ptr(z, 0, 0, 0, 0, 0), 6 * D,
        1.0, dRx_ptr(z, 0, 0, 0), 5 * D));
  }

  // dJ/dRy
  for (int z = 0; z < 4; ++z) {
    CHECK_CUBLAS_CALL(cublasSetStream(handle, stream[12 + z]));
    CHECK_CUBLAS_CALL(gemm_gpu<T>(
        handle, CUBLAS_OP_T, CUBLAS_OP_N, D, 5 * D, H * W * N,
        scale, Oyp + z * H * W * N * D, D,
        dQ_ptr(z, 0, 0, 0, 0, 0), 6 * D,
        1.0, dRy_ptr(z, 0, 0, 0), 5 * D));
  }

  STREAMS_SYNCHRONIZE(4 * 4);
