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

1D, 2D and 3D transpose cuda

Browse source
This commit is contained in:
lucasdelimanogueira 2024-05-23 02:09:32 -03:00
parent fb6429aad6
commit 702ebaf87c
5 changed files with 123 additions and 17 deletions
Download patch file Download diff file Expand all files Collapse all files

6
norch/csrc/cpu.cpp
View file

@ -375,14 +375,14 @@ void transpose_2D_tensor_cpu(Tensor* tensor, float* result_data) {
}
void transpose_3D_tensor_cpu(Tensor* tensor, float* result_data) {
int depth = tensor->shape[0];
int batch = tensor->shape[0];
int rows = tensor->shape[1];
int cols = tensor->shape[2];
for (int i = 0; i < depth; i++) {
for (int i = 0; i < batch; i++) {
for (int j = 0; j < rows; j++) {
for (int k = 0; k < cols; k++) {
result_data[k * rows * depth + j * depth + i] = tensor->data[i * rows * cols + j * cols + k];
result_data[k * rows * batch + j * batch + i] = tensor->data[i * rows * cols + j * cols + k];
}
}
}