 
   // Define device arrays and pass data from the CPU to GPU
   double *A, *B, *RetMat;
-  CUDA_CHECK(cudaMalloc((void**) &A, Q*NUM_THREADS*sizeof(double)));
-  CUDA_CHECK(cudaMalloc((void**) &B, Q*NUM_THREADS*sizeof(double)));
+  CUDA_CHECK(cudaMalloc((void**) &A, D*sizeof(double)));
+  CUDA_CHECK(cudaMalloc((void**) &B, Q*D*sizeof(double)));
   CUDA_CHECK(cudaMalloc((void**) &RetMat, Q*NUM_THREADS*sizeof(double)));
 
-  CUDA_CHECK(cudaMemcpy(A, prepA, Q*NUM_THREADS*sizeof(double), 
+  CUDA_CHECK(cudaMemcpy(A, X, D*sizeof(double), 
                         cudaMemcpyHostToDevice));
-  CUDA_CHECK(cudaMemcpy(B, prepB, Q*NUM_THREADS*sizeof(double), 
+  CUDA_CHECK(cudaMemcpy(B, Y, Q*D*sizeof(double), 
                         cudaMemcpyHostToDevice));
 
 /*

 */
/*
  printf("--- Print matrix A ---\n");
  print_GPU_Mat(A,Q*NUM_THREADS);
  printf("--- Print matrix B ---\n");
  print_GPU_Mat(B,Q*NUM_THREADS);
*//*
  printf("--- Print matrix RetMat ---\n");
  print_GPU_Mat(RetMat,D);

  // Define block and grid size
  const int y_dim = MAX_THREADS/NUM_THREADS;
  dim3 blockSize(NUM_THREADS,y_dim);
  const int num_blocks = Q/y_dim + 1;
  dim3 gridSize(num_blocks,1);
/*
  printf("Block size = %d, grid size = %d, number of elements (N) = %d\n", 
          NUM_THREADS, num_blocks, D);
*/
/*  // Matrix to prepare the data to pass to device
  double *prepA, *prepB;
  prepA = (double*)malloc(NUM_THREADS*Q*sizeof(double));
  prepB = (double*)malloc(NUM_THREADS*Q*sizeof(double));
