aparler/PyNorch
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

small fix 3D transpose

Browse source
This commit is contained in:
lucasdelimanogueira 2024-05-23 02:34:38 -03:00
parent 38e629fc38
commit a0bfd6a0eb
2 changed files with 2 additions and 2 deletions
Download patch file Download diff file Expand all files Collapse all files

2
norch/csrc/cuda.cu
View file

@ -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<<<number_of_blocks, threadsPerBlock>>>(tensor->data, result_data, rows, cols);
transpose_3D_tensor_cuda_kernel<<<number_of_blocks, threadsPerBlock>>>(tensor->data, result_data, batch, rows, cols);
cudaError_t error = cudaGetLastError();

2
norch/csrc/cuda.h
View file

@ -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);