From a0bfd6a0eb0fc8dc8725561962068337fd739f1b Mon Sep 17 00:00:00 2001 From: lucasdelimanogueira Date: Thu, 23 May 2024 02:34:38 -0300 Subject: [PATCH] small fix 3D transpose --- norch/csrc/cuda.cu | 2 +- norch/csrc/cuda.h | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/norch/csrc/cuda.cu b/norch/csrc/cuda.cu index 064c3d1..87ba247 100644 --- a/norch/csrc/cuda.cu +++ b/norch/csrc/cuda.cu @@ -1080,7 +1080,7 @@ __host__ void transpose_3D_tensor_cuda(Tensor* tensor, float* result_data) { dim3 threadsPerBlock(8, 8, 8); dim3 number_of_blocks((batch + threadsPerBlock.x - 1) / threadsPerBlock.x, (rows + threadsPerBlock.y - 1) / threadsPerBlock.y, (cols + threadsPerBlock.z - 1) / threadsPerBlock.z); - transpose_2D_tensor_cuda_kernel<<>>(tensor->data, result_data, rows, cols); + transpose_3D_tensor_cuda_kernel<<>>(tensor->data, result_data, batch, rows, cols); cudaError_t error = cudaGetLastError(); diff --git a/norch/csrc/cuda.h b/norch/csrc/cuda.h index cc7b385..daeef28 100644 --- a/norch/csrc/cuda.h +++ b/norch/csrc/cuda.h @@ -79,7 +79,7 @@ __global__ void transpose_2D_tensor_cuda_kernel(float* data, float* result_data, int rows, int cols); __host__ void transpose_2D_tensor_cuda(Tensor* tensor, float* result_data); - __global__ void transpose_3D_tensor_cuda_kernel(float* data, float* result_data, int rows, int cols); + __global__ void transpose_3D_tensor_cuda_kernel(float* data, float* result_data, int batch, int rows, int cols); __host__ void transpose_3D_tensor_cuda(Tensor* tensor, float* result_data); __global__ void assign_tensor_cuda_kernel(float* data, float* result_data, int size);