diff --git a/build/cpu.o b/build/cpu.o
index 2ee4a6f..ecb6dd6 100644
Binary files a/build/cpu.o and b/build/cpu.o differ
diff --git a/build/cuda.cu.o b/build/cuda.cu.o
index 4ed5a3e..4d96c07 100644
Binary files a/build/cuda.cu.o and b/build/cuda.cu.o differ
diff --git a/build/tensor.o b/build/tensor.o
index 47e5749..ee245d0 100644
Binary files a/build/tensor.o and b/build/tensor.o differ
diff --git a/norch/__pycache__/tensor.cpython-38.pyc b/norch/__pycache__/tensor.cpython-38.pyc
index 718f29e..bbd6f48 100644
Binary files a/norch/__pycache__/tensor.cpython-38.pyc and b/norch/__pycache__/tensor.cpython-38.pyc differ
diff --git a/norch/csrc/cpu.cpp b/norch/csrc/cpu.cpp
index 4e27010..63e1615 100644
--- a/norch/csrc/cpu.cpp
+++ b/norch/csrc/cpu.cpp
@@ -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];
             }
         }
     }
@@ -395,6 +395,25 @@ void assign_tensor_cpu(Tensor* tensor, float* result_data) {
     }
 }
 
+void make_contiguous_tensor_cpu(Tensor* tensor, float* result_data, int* new_strides) {
+    
+    for (int i = 0; i < tensor->size; i++) {
+        int index = 0;
+        int offset = i;
+        for (int j = 0; j < tensor->ndim; j++) {
+            index += (offset / new_strides[j]) * tensor->strides[j];
+            offset %= new_strides[j];
+        }
+        result_data[i] = tensor->data[index];
+    }
+
+    // Free old data and update tensor properties
+    free(tensor->data);
+    free(tensor->strides);
+    tensor->data = result_data;
+    tensor->strides = new_strides;
+}
+
 void sin_tensor_cpu(Tensor* tensor, float* result_data) {
     for (int i = 0; i < tensor->size; i++) {
         result_data[i] = sinf(tensor->data[i]);
diff --git a/norch/csrc/cpu.h b/norch/csrc/cpu.h
index 1d9547b..9a27a26 100644
--- a/norch/csrc/cpu.h
+++ b/norch/csrc/cpu.h
@@ -30,6 +30,7 @@ void transpose_2D_tensor_cpu(Tensor* tensor, float* result_data);
 void transpose_3D_tensor_cpu(Tensor* tensor, float* result_data);
 void transpose_axes_cpu(Tensor* tensor, float* result_data, int axis1, int axis2, int* new_shape);
 void assign_tensor_cpu(Tensor* tensor, float* result_data);
+void make_contiguous_tensor_cpu(Tensor* tensor, float* result_data, int* new_strides);
 void sin_tensor_cpu(Tensor* tensor, float* result_data);
 void cos_tensor_cpu(Tensor* tensor, float* result_data);
 void sigmoid_tensor_cpu(Tensor* tensor, float* result_data);
diff --git a/norch/csrc/cuda.cu b/norch/csrc/cuda.cu
index 6908780..87ba247 100644
--- a/norch/csrc/cuda.cu
+++ b/norch/csrc/cuda.cu
@@ -2,6 +2,7 @@
 #include <stdio.h>
 #include <stdlib.h>
 #include <math.h>
+#include <float.h>
 
 #define THREADS_PER_BLOCK 128
 #define TILE_SIZE 32
@@ -141,18 +142,20 @@ __global__ void sum_tensor_cuda_kernel(float* data, float* result_data, int size
     }
 }
 
-__global__ void sum_tensor_cuda_kernel_axis(float* data, float* result_data, int* strides, int target_axis, int inner_size, int outer_size) {
+__global__ void sum_tensor_cuda_kernel_axis(float* data, float* result_data, int* strides, int* shape, int axis, int ndim, int axis_stride, int size, int result_size) {
     int tid = blockIdx.x * blockDim.x + threadIdx.x;
 
-    if (tid < outer_size) {
-        int outer_index = tid / inner_size;
-        int inner_index = tid % inner_size;
+    if (tid < result_size) {
+        for (int i = 0; i < shape[axis]; i++) {
+            int index = 0;
+            int remainder = tid;
+            for (int k = ndim - 2; k >= 0; k--) {
+                index += (remainder % shape[k < axis ? k : k + 1]) * strides[k < axis ? k : k + 1];
+                remainder /= shape[k < axis ? k : k + 1];
+            }
+            index += i * axis_stride;
 
-        int offset = outer_index * strides[0] + inner_index;
-
-        for (int i = 0; i < strides[target_axis]; ++i) {
-            int index = offset + i * strides[target_axis + 1];
-            atomicAdd(&result_data[outer_index * inner_size + inner_index], data[index]);
+            atomicAdd(&result_data[tid], data[index]);
         }
     }
 }
@@ -184,20 +187,30 @@ __host__ void sum_tensor_cuda(Tensor* tensor, float* result_data, int axis) {
         cudaDeviceSynchronize();
         
     } else {
+        int axis_stride = tensor->strides[axis];
 
-        int target_axis_stride = tensor->strides[axis];
-        int inner_size = tensor->strides[axis + 1];
-        int outer_size = tensor->size / target_axis_stride;
+        // Calculate the size of the resulting tensor
+        int result_size = 1;
+        for (int i = 0; i < tensor->ndim; i++) {
+            if (i != axis) {
+                result_size *= tensor->shape[i];
+            }
+        }
 
+        // Allocate memory for strides and shape on the device
         int* d_strides;
-        cudaMalloc(&d_strides, (tensor->ndim + 1) * sizeof(int));
-        cudaMemcpy(d_strides, tensor->strides, (tensor->ndim + 1) * sizeof(int), cudaMemcpyHostToDevice);
+        int* d_shape;
+        cudaMalloc(&d_strides, tensor->ndim * sizeof(int));
+        cudaMalloc(&d_shape, tensor->ndim * sizeof(int));
+        cudaMemcpy(d_strides, tensor->strides, tensor->ndim * sizeof(int), cudaMemcpyHostToDevice);
+        cudaMemcpy(d_shape, tensor->shape, tensor->ndim * sizeof(int), cudaMemcpyHostToDevice);
 
-        cudaMemset(result_data, 0, outer_size * sizeof(float));
+        // Initialize result_data to 0
+        cudaMemset(result_data, 0, result_size * sizeof(float));
 
-        int num_threads = outer_size;
+        int num_threads = result_size;
         int num_blocks = (num_threads + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK;
-        sum_tensor_cuda_kernel_axis<<<num_blocks, THREADS_PER_BLOCK>>>(tensor->data, result_data, d_strides, axis, inner_size, outer_size);
+        sum_tensor_cuda_kernel_axis<<<num_blocks, THREADS_PER_BLOCK>>>(tensor->data, result_data, d_strides, d_shape, axis, tensor->ndim, axis_stride, tensor->size, result_size);
 
         cudaError_t error = cudaGetLastError();
         if (error != cudaSuccess) {
@@ -209,9 +222,254 @@ __host__ void sum_tensor_cuda(Tensor* tensor, float* result_data, int axis) {
 
         // Free allocated memory
         cudaFree(d_strides);
+        cudaFree(d_shape);
     }
 }
 
+__global__ void max_tensor_cuda_kernel(float* data, float* result_data, int size) {
+    __shared__ float partial_max[THREADS_PER_BLOCK * sizeof(float)];
+
+    int tid = threadIdx.x;
+    int i = blockIdx.x * blockDim.x + threadIdx.x;
+
+    partial_max[tid] = (i < size) ? data[i] : -FLT_MAX;
+
+    __syncthreads();
+
+    // Perform block-wise reduction
+    for (int s = blockDim.x / 2; s > 0; s >>= 1) {
+        if (tid < s) {
+            partial_max[tid] = fmax(partial_max[tid], partial_max[tid + s]);
+        }
+        __syncthreads();
+    }
+
+    // Write block sum to global memory
+    if (tid == 0) {
+        result_data[blockIdx.x] = partial_max[0];
+    }
+}
+
+__device__ float atomicMaxFloat(float* address, float val) {
+    int* address_as_int = (int*)address;
+    int old = *address_as_int, assumed;
+
+    do {
+        assumed = old;
+        old = atomicCAS(address_as_int, assumed,
+                        __float_as_int(fmaxf(val, __int_as_float(assumed))));
+    } while (assumed != old);
+
+    return __int_as_float(old);
+}
+
+__global__ void max_tensor_cuda_kernel_axis(float* data, float* result_data, int* strides, int* shape, int axis, int ndim, int axis_stride, int size, int result_size) {
+    int tid = blockIdx.x * blockDim.x + threadIdx.x;
+
+    if (tid < result_size) {
+        for (int i = 0; i < shape[axis]; i++) {
+            int index = 0;
+            int remainder = tid;
+            for (int k = ndim - 2; k >= 0; k--) {
+                index += (remainder % shape[k < axis ? k : k + 1]) * strides[k < axis ? k : k + 1];
+                remainder /= shape[k < axis ? k : k + 1];
+            }
+            index += i * axis_stride;
+
+            atomicMaxFloat(&result_data[tid], data[index]);
+        }
+    }
+}
+
+__host__ void max_tensor_cuda(Tensor* tensor, float* result_data, int axis) {
+
+    if (axis == -1) {
+        cudaMemcpy(result_data, tensor->data, tensor->size * sizeof(float), cudaMemcpyHostToDevice);
+        
+        int num_blocks = (tensor->size + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK;
+
+        // First-level reduction
+        max_tensor_cuda_kernel<<<num_blocks, THREADS_PER_BLOCK>>>(tensor->data, result_data, tensor->size);
+
+        // If necessary, perform multiple levels of reduction
+        while (num_blocks > 1) {
+            int num_blocks_next = (num_blocks + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK;
+            max_tensor_cuda_kernel<<<num_blocks_next, THREADS_PER_BLOCK>>>(result_data, result_data, num_blocks);
+            num_blocks = num_blocks_next;
+        }
+
+        cudaError_t error = cudaGetLastError();
+        if (error != cudaSuccess) {
+            printf("CUDA error: %s\n", cudaGetErrorString(error));
+            exit(-1);
+        }
+
+        cudaDeviceSynchronize();
+        
+    } else {
+        int axis_stride = tensor->strides[axis];
+
+        // Calculate the size of the resulting tensor
+        int result_size = 1;
+        for (int i = 0; i < tensor->ndim; i++) {
+            if (i != axis) {
+                result_size *= tensor->shape[i];
+            }
+        }
+
+        // Allocate memory for strides and shape on the device
+        int* d_strides;
+        int* d_shape;
+        cudaMalloc(&d_strides, tensor->ndim * sizeof(int));
+        cudaMalloc(&d_shape, tensor->ndim * sizeof(int));
+        cudaMemcpy(d_strides, tensor->strides, tensor->ndim * sizeof(int), cudaMemcpyHostToDevice);
+        cudaMemcpy(d_shape, tensor->shape, tensor->ndim * sizeof(int), cudaMemcpyHostToDevice);
+
+        // Initialize result_data to 0
+        float neg_inf = -FLT_MAX;
+        cudaMemset(result_data, *reinterpret_cast<int*>(&neg_inf), result_size * sizeof(float));
+
+        int num_threads = result_size;
+        int num_blocks = (num_threads + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK;
+        max_tensor_cuda_kernel_axis<<<num_blocks, THREADS_PER_BLOCK>>>(tensor->data, result_data, d_strides, d_shape, axis, tensor->ndim, axis_stride, tensor->size, result_size);
+
+        cudaError_t error = cudaGetLastError();
+        if (error != cudaSuccess) {
+            printf("CUDA error: %s\n", cudaGetErrorString(error));
+            exit(-1);
+        }
+
+        cudaDeviceSynchronize();
+
+        // Free allocated memory
+        cudaFree(d_strides);
+        cudaFree(d_shape);
+    }
+}
+
+__global__ void min_tensor_cuda_kernel(float* data, float* result_data, int size) {
+    __shared__ float partial_min[THREADS_PER_BLOCK * sizeof(float)];
+
+    int tid = threadIdx.x;
+    int i = blockIdx.x * blockDim.x + threadIdx.x;
+
+    partial_min[tid] = (i < size) ? data[i] : FLT_MAX;
+
+    __syncthreads();
+
+    // Perform block-wise reduction
+    for (int s = blockDim.x / 2; s > 0; s >>= 1) {
+        if (tid < s) {
+            partial_min[tid] = fmin(partial_min[tid], partial_min[tid + s]);
+        }
+        __syncthreads();
+    }
+
+    // Write block sum to global memory
+    if (tid == 0) {
+        result_data[blockIdx.x] = partial_min[0];
+    }
+}
+
+__device__ float atomicMinFloat(float* address, float val) {
+    int* address_as_int = (int*)address;
+    int old = *address_as_int, assumed;
+
+    do {
+        assumed = old;
+        old = atomicCAS(address_as_int, assumed,
+                        __float_as_int(fminf(val, __int_as_float(assumed))));
+    } while (assumed != old);
+
+    return __int_as_float(old);
+}
+
+__global__ void min_tensor_cuda_kernel_axis(float* data, float* result_data, int* strides, int* shape, int axis, int ndim, int axis_stride, int size, int result_size) {
+    int tid = blockIdx.x * blockDim.x + threadIdx.x;
+
+    if (tid < result_size) {
+        for (int i = 0; i < shape[axis]; i++) {
+            int index = 0;
+            int remainder = tid;
+            for (int k = ndim - 2; k >= 0; k--) {
+                index += (remainder % shape[k < axis ? k : k + 1]) * strides[k < axis ? k : k + 1];
+                remainder /= shape[k < axis ? k : k + 1];
+            }
+            index += i * axis_stride;
+
+            atomicMinFloat(&result_data[tid], data[index]);
+        }
+    }
+}
+
+__host__ void min_tensor_cuda(Tensor* tensor, float* result_data, int axis) {
+
+    if (axis == -1) {
+        cudaMemcpy(result_data, tensor->data, tensor->size * sizeof(float), cudaMemcpyHostToDevice);
+        
+        int num_blocks = (tensor->size + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK;
+
+        // First-level reduction
+        min_tensor_cuda_kernel<<<num_blocks, THREADS_PER_BLOCK>>>(tensor->data, result_data, tensor->size);
+
+        // If necessary, perform multiple levels of reduction
+        while (num_blocks > 1) {
+            int num_blocks_next = (num_blocks + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK;
+            min_tensor_cuda_kernel<<<num_blocks_next, THREADS_PER_BLOCK>>>(result_data, result_data, num_blocks);
+            num_blocks = num_blocks_next;
+        }
+
+        cudaError_t error = cudaGetLastError();
+        if (error != cudaSuccess) {
+            printf("CUDA error: %s\n", cudaGetErrorString(error));
+            exit(-1);
+        }
+
+        cudaDeviceSynchronize();
+        
+    } else {
+        int axis_stride = tensor->strides[axis];
+
+        // Calculate the size of the resulting tensor
+        int result_size = 1;
+        for (int i = 0; i < tensor->ndim; i++) {
+            if (i != axis) {
+                result_size *= tensor->shape[i];
+            }
+        }
+
+        // Allocate memory for strides and shape on the device
+        int* d_strides;
+        int* d_shape;
+        cudaMalloc(&d_strides, tensor->ndim * sizeof(int));
+        cudaMalloc(&d_shape, tensor->ndim * sizeof(int));
+        cudaMemcpy(d_strides, tensor->strides, tensor->ndim * sizeof(int), cudaMemcpyHostToDevice);
+        cudaMemcpy(d_shape, tensor->shape, tensor->ndim * sizeof(int), cudaMemcpyHostToDevice);
+
+        // Initialize result_data to 0
+        float inf = FLT_MAX;
+        cudaMemset(result_data, *reinterpret_cast<int*>(&inf), result_size * sizeof(float));
+
+        int num_threads = result_size;
+        int num_blocks = (num_threads + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK;
+        min_tensor_cuda_kernel_axis<<<num_blocks, THREADS_PER_BLOCK>>>(tensor->data, result_data, d_strides, d_shape, axis, tensor->ndim, axis_stride, tensor->size, result_size);
+
+        cudaError_t error = cudaGetLastError();
+        if (error != cudaSuccess) {
+            printf("CUDA error: %s\n", cudaGetErrorString(error));
+            exit(-1);
+        }
+
+        cudaDeviceSynchronize();
+
+        // Free allocated memory
+        cudaFree(d_strides);
+        cudaFree(d_shape);
+    }
+}
+
+
+
 __global__ void sub_tensor_cuda_kernel(float* data1, float* data2, float* result_data, int size) {
    
     int i = blockIdx.x * blockDim.x + threadIdx.x;
@@ -753,23 +1011,46 @@ __host__ void zeros_like_tensor_cuda(Tensor* tensor, float* result_data) {
     cudaDeviceSynchronize();
 }
 
-__global__ void transpose_tensor_cuda_kernel(float* data, float* result_data, int rows, int cols) {
-    int tid_x = blockIdx.x * blockDim.x + threadIdx.x;
-    int tid_y = blockIdx.y * blockDim.y + threadIdx.y;
 
-    if (tid_x < cols && tid_y < rows) {
-        result_data[tid_x * rows + tid_y] = data[tid_y * cols + tid_x];
+__global__ void transpose_1D_tensor_cuda_kernel(float* data, float* result_data, int size) {
+    
+    int i = blockIdx.x * blockDim.x + threadIdx.x;
+    if (i < size) {
+        result_data[i] = data[i];
     }
 }
 
-__host__ void transpose_tensor_cuda(Tensor* tensor, float* result_data) {
+__host__ void transpose_1D_tensor_cuda(Tensor* tensor, float* result_data) {
+    
+    int number_of_blocks = (tensor->size + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK;
+    transpose_1D_tensor_cuda_kernel<<<number_of_blocks, THREADS_PER_BLOCK>>>(tensor->data, result_data, tensor->size);
+
+    cudaError_t error = cudaGetLastError();
+    if (error != cudaSuccess) {
+        printf("CUDA error: %s\n", cudaGetErrorString(error));
+        exit(-1);
+    }
+
+    cudaDeviceSynchronize();
+}
+
+__global__ void transpose_2D_tensor_cuda_kernel(float* data, float* result_data, int rows, int cols) {
+    int i = blockIdx.x * blockDim.x + threadIdx.x;
+    int j = blockIdx.y * blockDim.y + threadIdx.y;
+
+    if (i < rows && j < cols) {
+        result_data[j * rows + i] = data[i * cols + j];
+    }
+}
+
+__host__ void transpose_2D_tensor_cuda(Tensor* tensor, float* result_data) {
     
     int rows = tensor->shape[0];
     int cols = tensor->shape[1];
 
     dim3 threadsPerBlock(16, 16);
-    dim3 number_of_blocks((cols + threadsPerBlock.x - 1) / threadsPerBlock.x, (rows + threadsPerBlock.y - 1) / threadsPerBlock.y);
-    transpose_tensor_cuda_kernel<<<number_of_blocks, threadsPerBlock>>>(tensor->data, result_data, rows, cols);
+    dim3 number_of_blocks((rows + threadsPerBlock.x - 1) / threadsPerBlock.x, (cols + threadsPerBlock.y - 1) / threadsPerBlock.y);
+    transpose_2D_tensor_cuda_kernel<<<number_of_blocks, threadsPerBlock>>>(tensor->data, result_data, rows, cols);
 
 
     cudaError_t error = cudaGetLastError();
@@ -781,6 +1062,37 @@ __host__ void transpose_tensor_cuda(Tensor* tensor, float* result_data) {
     cudaDeviceSynchronize();
 }
 
+__global__ void transpose_3D_tensor_cuda_kernel(float* data, float* result_data, int batch, int rows, int cols) {
+    int i = blockIdx.x * blockDim.x + threadIdx.x;
+    int j = blockIdx.y * blockDim.y + threadIdx.y;
+    int k = blockIdx.z * blockDim.z + threadIdx.z;
+
+    if (i < batch && j < rows && k < cols) {
+        result_data[k * rows * batch + j * batch + i] = data[i * rows * cols + j * cols + k];
+    }
+}
+
+__host__ void transpose_3D_tensor_cuda(Tensor* tensor, float* result_data) {
+    
+    int batch = tensor->shape[0];
+    int rows = tensor->shape[1];
+    int cols = tensor->shape[2];
+
+    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_3D_tensor_cuda_kernel<<<number_of_blocks, threadsPerBlock>>>(tensor->data, result_data, batch, rows, cols);
+
+
+    cudaError_t error = cudaGetLastError();
+    if (error != cudaSuccess) {
+        printf("CUDA error: %s\n", cudaGetErrorString(error));
+        exit(-1);
+    }
+
+    cudaDeviceSynchronize();
+}
+
+
 __global__ void assign_tensor_cuda_kernel(float* data, float* result_data, int size) {
     
     int i = blockIdx.x * blockDim.x + threadIdx.x;
@@ -803,6 +1115,48 @@ __host__ void assign_tensor_cuda(Tensor* tensor, float* result_data) {
     cudaDeviceSynchronize();
 }
 
+__global__ void make_contiguous_tensor_cuda_kernel(float* data, float* result_data, int ndim, int size, int* strides, int* new_strides) {
+    
+    int i = blockIdx.x * blockDim.x + threadIdx.x;
+    if (i < size) {
+        int index = 0;
+        int offset = i;
+        for (int j = 0; j < ndim; j++) {
+            index += (offset / new_strides[j]) * strides[j];
+            offset %= new_strides[j];
+        }
+        result_data[i] = data[index];
+    }
+}
+
+__host__ void make_contiguous_tensor_cuda(Tensor* tensor, float* result_data, int* new_strides) {
+    
+    int* d_strides;
+    cudaMalloc((void **)&d_strides, tensor->ndim * sizeof(int));
+    cudaMemcpy(d_strides, tensor->strides, tensor->ndim * sizeof(int), cudaMemcpyHostToDevice);
+    
+    int* d_new_strides;
+    cudaMalloc((void **)&d_new_strides, tensor->ndim * sizeof(int));
+    cudaMemcpy(d_new_strides, new_strides, tensor->ndim * sizeof(int), cudaMemcpyHostToDevice);
+
+    int number_of_blocks = (tensor->size + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK;
+    make_contiguous_tensor_cuda_kernel<<<number_of_blocks, THREADS_PER_BLOCK>>>(tensor->data, result_data, tensor->ndim, tensor->size, d_strides, d_new_strides);
+
+    cudaError_t error = cudaGetLastError();
+    if (error != cudaSuccess) {
+        printf("CUDA error: %s\n", cudaGetErrorString(error));
+        exit(-1);
+    }
+
+    cudaDeviceSynchronize();
+
+    // Free old data and update tensor properties
+    cudaFree(tensor->data);
+    free(tensor->strides);
+    tensor->data = result_data;
+    tensor->strides = new_strides;
+}
+
 __global__ void sin_tensor_cuda_kernel(float* data, float* result_data, int size) {
     
     int i = blockIdx.x * blockDim.x + threadIdx.x;
diff --git a/norch/csrc/cuda.h b/norch/csrc/cuda.h
index 822d661..daeef28 100644
--- a/norch/csrc/cuda.h
+++ b/norch/csrc/cuda.h
@@ -14,9 +14,17 @@
     __host__ void sub_broadcasted_tensor_cuda(Tensor* tensor1, Tensor* tensor2, float* result_data, int* broadcasted_shape, int broadcasted_size);
 
     __global__ void sum_tensor_cuda_kernel(float* data, float* result_data);
-    __global__ void sum_tensor_cuda_kernel_axis(float* data, float* result_data, int* strides, int target_axis, int inner_size, int outer_size);
+    __global__ void sum_tensor_cuda_kernel_axis(float* data, float* result_data, int* strides, int* shape, int axis, int ndim, int axis_stride, int size, int result_size);
     __host__ void sum_tensor_cuda(Tensor* tensor, float* result_data, int axis);
 
+    __global__ void max_tensor_cuda_kernel(float* data, float* result_data);
+    __global__ void max_tensor_cuda_kernel_axis(float* data, float* result_data, int* strides, int* shape, int axis, int ndim, int axis_stride, int size, int result_size);
+    __host__ void max_tensor_cuda(Tensor* tensor, float* result_data, int axis);
+
+    __global__ void min_tensor_cuda_kernel(float* data, float* result_data);
+    __global__ void min_tensor_cuda_kernel_axis(float* data, float* result_data, int* strides, int* shape, int axis, int ndim, int axis_stride, int size, int result_size);
+    __host__ void min_tensor_cuda(Tensor* tensor, float* result_data, int axis);
+
     __global__ void sub_tensor_cuda_kernel(float* data1, float* data2, float* result_data, int size);
     __host__ void sub_tensor_cuda(Tensor* tensor1, Tensor* tensor2, float* result_data);
     
@@ -65,12 +73,21 @@
     __global__ void zeros_like_tensor_cuda_kernel(float* data, float* result_data, int size);
     __host__ void zeros_like_tensor_cuda(Tensor* tensor, float* result_data);
 
-    __global__ void transpose_tensor_cuda_kernel(float* data, float* result_data, int rows, int cols);
-    __host__ void transpose_tensor_cuda(Tensor* tensor, float* result_data);
+    __global__ void transpose_1D_tensor_cuda_kernel(float* data, float* result_data, int rows, int cols);
+    __host__ void transpose_1D_tensor_cuda(Tensor* tensor, float* result_data);
+
+    __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 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);
     __host__ void assign_tensor_cuda(Tensor* tensor, float* result_data);
 
+    __global__ void make_contiguous_tensor_cuda_kernel(float* data, float* result_data, int ndim, int size, int* strides, int* new_strides);
+    __host__ void make_contiguous_tensor_cuda(Tensor* tensor, float* result_data, int* new_strides);
+
     __global__ void sin_tensor_cuda_kernel(float* data, float* result_data, int size);
     __host__ void sin_tensor_cuda(Tensor* tensor, float* result_data);
 
diff --git a/norch/csrc/tensor.cpp b/norch/csrc/tensor.cpp
index 221ec62..d5bfa71 100644
--- a/norch/csrc/tensor.cpp
+++ b/norch/csrc/tensor.cpp
@@ -202,9 +202,9 @@ extern "C" {
             ndim = tensor->ndim - 1;
         }
 
-        int size = 1;
+        int axis_size = 1;
         for (int i = 0; i < ndim; i++) {
-            size *= shape[i];
+            axis_size *= shape[i];
         }
   
         if (strcmp(tensor->device, "cuda") == 0) {
@@ -213,11 +213,7 @@ extern "C" {
             if (axis == -1) {
                 cudaMalloc((void**)&result_data, tensor->size * sizeof(float));
             } else {
-
-                int target_axis_stride = tensor->strides[axis];
-                int outer_size = tensor->size / target_axis_stride;
-
-                cudaMalloc((void**)&result_data, (outer_size) * sizeof(float));
+                cudaMalloc((void**)&result_data, axis_size * sizeof(float));
             }
             sum_tensor_cuda(tensor, result_data, axis);
             
@@ -241,13 +237,13 @@ extern "C" {
             return create_tensor(result_data, shape, ndim, device);
         } 
         else {
-            float* result_data = (float*)calloc(size, sizeof(float));
+            float* result_data = (float*)calloc(axis_size, sizeof(float));
             if (result_data == NULL) {
                 fprintf(stderr, "Memory allocation failed\n");
                 exit(1);
             }
 
-            sum_tensor_cpu(tensor, result_data, size, shape, axis);
+            sum_tensor_cpu(tensor, result_data, axis_size, shape, axis);
 
             if (keepdim) {
                 if (axis == -1){
@@ -294,26 +290,20 @@ extern "C" {
             ndim = tensor->ndim - 1;
         }
 
-        int size = 1;
+        int axis_size = 1;
         for (int i = 0; i < ndim; i++) {
-            size *= shape[i];
+            axis_size *= shape[i];
         }
   
         if (strcmp(tensor->device, "cuda") == 0) {
 
             float* result_data;
-            cudaMalloc((void**)&result_data, size * sizeof(float));
-            //cudaMemset(result_data, -INFINITY, size * sizeof(float));
-            //max_tensor_cuda(tensor, result_data);
-            return create_tensor(result_data, shape, ndim, device);
-        } 
-        else {
-            float* result_data = (float*)malloc(size * sizeof(float));
-            if (result_data == NULL) {
-                fprintf(stderr, "Memory allocation failed\n");
-                exit(1);
+            if (axis == -1) {
+                cudaMalloc((void**)&result_data, tensor->size * sizeof(float));
+            } else {
+                cudaMalloc((void**)&result_data, axis_size * sizeof(float));
             }
-            max_tensor_cpu(tensor, result_data, size, shape, axis);
+            max_tensor_cuda(tensor, result_data, axis);
 
             if (keepdim) {
                 if (axis == -1){
@@ -329,9 +319,37 @@ extern "C" {
                     }
                     shape[axis] = 1;
                     ndim = tensor->ndim;
-                }   
+                }
+                
             }
+            return create_tensor(result_data, shape, ndim, device);
 
+        } 
+        else {
+            float* result_data = (float*)malloc(axis_size * sizeof(float));
+            if (result_data == NULL) {
+                fprintf(stderr, "Memory allocation failed\n");
+                exit(1);
+            }
+            max_tensor_cpu(tensor, result_data, axis_size, shape, axis);
+
+            if (keepdim) {
+                if (axis == -1){
+                    ndim = tensor->ndim;
+                    shape = (int*) malloc((tensor->ndim) * sizeof(int));
+                    for (int i = 0; i < tensor->ndim; i++) {
+                        shape[i] = 1;
+                    }
+                } else {
+                    shape = (int*) malloc((tensor->ndim) * sizeof(int));
+                    for (int i = 0; i < tensor->ndim; i++) {
+                        shape[i] = tensor->shape[i];
+                    }
+                    shape[axis] = 1;
+                    ndim = tensor->ndim;
+                }
+                
+            }
             return create_tensor(result_data, shape, ndim, device);
         }     
     }
@@ -347,7 +365,6 @@ extern "C" {
         int ndim;
         int* shape;
         if (axis == -1) {
-            
             shape = (int*) malloc(sizeof(int));
             shape[0] = 1;
             ndim = 1;
@@ -361,26 +378,20 @@ extern "C" {
             ndim = tensor->ndim - 1;
         }
         
-        int size = 1;
+        int axis_size = 1;
         for (int i = 0; i < ndim; i++) {
-            size *= shape[i];
+            axis_size *= shape[i];
         }
   
         if (strcmp(tensor->device, "cuda") == 0) {
 
             float* result_data;
-            cudaMalloc((void**)&result_data, size * sizeof(float));
-            //cudaMemset(result_data, INFINITY, size * sizeof(float));
-            //min_tensor_cuda(tensor, result_data);
-            return create_tensor(result_data, shape, ndim, device);
-        } 
-        else {
-            float* result_data = (float*)malloc(size * sizeof(float));
-            if (result_data == NULL) {
-                fprintf(stderr, "Memory allocation failed\n");
-                exit(1);
+            if (axis == -1) {
+                cudaMalloc((void**)&result_data, tensor->size * sizeof(float));
+            } else {
+                cudaMalloc((void**)&result_data, axis_size * sizeof(float));
             }
-            min_tensor_cpu(tensor, result_data, size, shape, axis);
+            min_tensor_cuda(tensor, result_data, axis);
 
             if (keepdim) {
                 if (axis == -1){
@@ -396,11 +407,39 @@ extern "C" {
                     }
                     shape[axis] = 1;
                     ndim = tensor->ndim;
-                }   
+                }
+                
             }
-
             return create_tensor(result_data, shape, ndim, device);
-        }     
+
+        } 
+        else {
+            float* result_data = (float*)malloc(axis_size * sizeof(float));
+            if (result_data == NULL) {
+                fprintf(stderr, "Memory allocation failed\n");
+                exit(1);
+            }
+            min_tensor_cpu(tensor, result_data, axis_size, shape, axis);
+
+            if (keepdim) {
+                if (axis == -1){
+                    ndim = tensor->ndim;
+                    shape = (int*) malloc((tensor->ndim) * sizeof(int));
+                    for (int i = 0; i < tensor->ndim; i++) {
+                        shape[i] = 1;
+                    }
+                } else {
+                    shape = (int*) malloc((tensor->ndim) * sizeof(int));
+                    for (int i = 0; i < tensor->ndim; i++) {
+                        shape[i] = tensor->shape[i];
+                    }
+                    shape[axis] = 1;
+                    ndim = tensor->ndim;
+                }
+                
+            }
+            return create_tensor(result_data, shape, ndim, device);
+        }       
     }
 
     Tensor* sub_tensor(Tensor* tensor1, Tensor* tensor2) {
@@ -1383,7 +1422,20 @@ extern "C" {
 
             float* result_data;
             cudaMalloc((void **)&result_data, size * sizeof(float));
-            transpose_tensor_cuda(tensor, result_data);
+            switch (ndim) {
+                case 1:
+                    transpose_1D_tensor_cuda(tensor, result_data);
+                    break;
+                case 2:
+                    transpose_2D_tensor_cuda(tensor, result_data);
+                    break;
+                case 3:
+                    transpose_3D_tensor_cuda(tensor, result_data);
+                    break;
+                default:
+                    fprintf(stderr, "Transpose only supports tensors up to 3 dimensions.\n");
+                    exit(-1);
+            }
             return create_tensor(result_data, shape, ndim, device);
         } 
         else {
@@ -1440,7 +1492,15 @@ extern "C" {
             float* result_data;
             cudaMalloc((void **)&result_data, size * sizeof(float));
             assign_tensor_cuda(tensor, result_data);
-            return create_tensor(result_data, shape, ndim, device);
+
+            Tensor* new_tensor = create_tensor(result_data, shape, ndim, device);
+            for (int i = 0; i < ndim; i++) {
+                new_tensor->strides[i] = tensor->strides[i];
+            }
+            new_tensor->strides[axis1] = tensor->strides[axis2];
+            new_tensor->strides[axis2] = tensor->strides[axis1];
+            make_contiguous(new_tensor);
+            return new_tensor;
         } 
         else {
             float* result_data = (float*)malloc(size * sizeof(float));
@@ -1462,17 +1522,10 @@ extern "C" {
     }
 
     void make_contiguous(Tensor* tensor) {
-        float* new_data = (float*)malloc(tensor->size * sizeof(float));
-        if (new_data == NULL) {
-            // Handle memory allocation failure
-            return;
-        }
-
+        
         int* new_strides = (int*)malloc(tensor->ndim * sizeof(int));
         if (new_strides == NULL) {
-            free(new_data);
-            // Handle memory allocation failure
-            return;
+            fprintf(stderr, "Memory allocation failed\n");
         }
 
         // Calculate new strides assuming C-contiguous order
@@ -1482,21 +1535,17 @@ extern "C" {
             stride *= tensor->shape[i];
         }
 
-        // Rearrange data
-        for (int i = 0; i < tensor->size; i++) {
-            int index = 0;
-            int offset = i;
-            for (int j = 0; j < tensor->ndim; j++) {
-                index += (offset / new_strides[j]) * tensor->strides[j];
-                offset %= new_strides[j];
-            }
-            new_data[i] = tensor->data[index];
-        }
+        if (strcmp(tensor->device, "cuda") == 0) {
+            float* result_data;
+            cudaMalloc((void **)&result_data, tensor->size * sizeof(float));
+            make_contiguous_tensor_cuda(tensor, result_data, new_strides);
 
-        // Free old data and update tensor properties
-        free(tensor->data);
-        free(tensor->strides);
-        tensor->data = new_data;
-        tensor->strides = new_strides;
+        } else {
+            float* result_data = (float*)malloc(tensor->size * sizeof(float));
+            if (result_data == NULL) {
+                fprintf(stderr, "Memory allocation failed\n");
+            }
+            make_contiguous_tensor_cpu(tensor, result_data, new_strides);
         }
     }
+}
diff --git a/norch/libtensor.so b/norch/libtensor.so
index eb4a5ac..7ce5299 100755
Binary files a/norch/libtensor.so and b/norch/libtensor.so differ
diff --git a/test.py b/test.py
new file mode 100644
index 0000000..94f5195
--- /dev/null
+++ b/test.py
@@ -0,0 +1,7 @@
+import norch
+from norch.utils import utils_unittests as utils
+
+device = "cpu"
+
+norch_tensor = norch.Tensor([[[[1, 2], [3, -4]], [[5, 6], [7, 8]]], [[[1, 2], [3, -4]], [[5, 6], [7, 8]]]]).to(device)
+norch_result = norch_tensor.sum(axis=0)
diff --git a/tests/test_autograd.py b/tests/test_autograd.py
index 2265c6d..82e6b61 100644
--- a/tests/test_autograd.py
+++ b/tests/test_autograd.py
@@ -24,8 +24,12 @@ class TestTensorAutograd(unittest.TestCase):
         norch_tensor1_grad = utils.to_torch(norch_tensor1.grad).to(self.device)
         norch_tensor2_grad = utils.to_torch(norch_tensor2.grad).to(self.device)
 
-        torch_tensor1 = torch.tensor([[[1, 2.5], [3, -4]], [[5, 6], [7, 8]]], requires_grad=True).to(self.device)
-        torch_tensor2 = torch.tensor([[[1, 1.], [1, 1.9]], [[1, 1], [1, 1]]], requires_grad=True).to(self.device)
+        torch_tensor1 = torch.tensor([[[1, 2.5], [3, -4]], [[5, 6], [7, 8]]], requires_grad=True)
+        torch_tensor2 = torch.tensor([[[1, 1.], [1, 1.9]], [[1, 1], [1, 1]]], requires_grad=True)
+
+        torch_tensor1.to(self.device)
+        torch_tensor2.to(self.device)
+
         torch_result = (torch_tensor1 + torch_tensor2).sum()
         torch_result.backward()
         torch_tensor1_grad = torch_tensor1.grad
@@ -46,8 +50,11 @@ class TestTensorAutograd(unittest.TestCase):
         norch_tensor1_grad = utils.to_torch(norch_tensor1.grad).to(self.device)
         norch_tensor2_grad = utils.to_torch(norch_tensor2.grad).to(self.device)
 
-        torch_tensor1 = torch.tensor([[[1, 2.5], [3, -4]], [[5, 6], [7, 8]]], requires_grad=True).to(self.device)
-        torch_tensor2 = torch.tensor([[[1, 1.], [1, 1.9]], [[1, 1], [1, 1]]], requires_grad=True).to(self.device)
+        torch_tensor1 = torch.tensor([[[1, 2.5], [3, -4]], [[5, 6], [7, 8]]], requires_grad=True)
+        torch_tensor2 = torch.tensor([[[1, 1.], [1, 1.9]], [[1, 1], [1, 1]]], requires_grad=True)
+
+        torch_tensor1.to(self.device)
+        torch_tensor2.to(self.device)
         
         torch_result = (torch_tensor1 + torch_tensor2).sum(axis=0).sum(axis=0).sum()
         torch_result.backward()
@@ -65,8 +72,11 @@ class TestTensorAutograd(unittest.TestCase):
         norch_tensor1_grad = utils.to_torch(norch_tensor1.grad).to(self.device)
         norch_tensor2_grad = utils.to_torch(norch_tensor2.grad).to(self.device)
 
-        torch_tensor1 = torch.tensor([[[1, 2.5], [3, -4]], [[5, 6], [7, 8]]], requires_grad=True).to(self.device)
-        torch_tensor2 = torch.tensor([[[1, 1.], [1, 1.9]], [[1, 1], [1, 1]]], requires_grad=True).to(self.device)
+        torch_tensor1 = torch.tensor([[[1, 2.5], [3, -4]], [[5, 6], [7, 8]]], requires_grad=True)
+        torch_tensor2 = torch.tensor([[[1, 1.], [1, 1.9]], [[1, 1], [1, 1]]], requires_grad=True)
+
+        torch_tensor1.to(self.device)
+        torch_tensor2.to(self.device)
         
         torch_result = (torch_tensor1 + torch_tensor2).sum(axis=1).sum()
         torch_result.backward()
@@ -87,7 +97,9 @@ class TestTensorAutograd(unittest.TestCase):
         norch_result.backward()
         norch_tensor_grad = utils.to_torch(norch_tensor.grad).to(self.device)
 
-        torch_tensor = torch.tensor([[[10, 10], [-4, -4]], [[5., 6], [7, 8]]], requires_grad=True).to(self.device)
+        torch_tensor = torch.tensor([[[10, 10], [-4, -4]], [[5., 6], [7, 8]]], requires_grad=True)
+
+        torch_tensor.to(self.device)
         
         torch_result = torch_tensor.max()
         torch_result.backward()
@@ -106,7 +118,9 @@ class TestTensorAutograd(unittest.TestCase):
         norch_result.backward()
         norch_tensor_grad = utils.to_torch(norch_tensor.grad).to(self.device)
 
-        torch_tensor = torch.tensor([[[10, 10], [-4, -4]], [[5., 6], [7, 8]]], requires_grad=True).to(self.device)
+        torch_tensor = torch.tensor([[[10, 10], [-4, -4]], [[5., 6], [7, 8]]], requires_grad=True)
+        
+        torch_tensor.to(self.device)
 
         torch_max_axis, _ = torch_tensor.max(axis=1)
         torch_result = torch_max_axis.sum()
@@ -122,7 +136,9 @@ class TestTensorAutograd(unittest.TestCase):
         norch_result.backward()
         norch_tensor_grad = utils.to_torch(norch_tensor.grad).to(self.device)
 
-        torch_tensor = torch.tensor([[[10, 1], [-4, 0]], [[5., 50], [7, 8]]], requires_grad=True).to(self.device)
+        torch_tensor = torch.tensor([[[10, 1], [-4, 0]], [[5., 50], [7, 8]]], requires_grad=True)
+
+        torch_tensor.to(self.device)
 
         torch_max_axis, _ = torch_tensor.max(axis=2)
         torch_result = torch_max_axis.sum()
@@ -186,7 +202,9 @@ class TestTensorAutograd(unittest.TestCase):
         norch_result.backward()
         norch_tensor_grad = utils.to_torch(norch_tensor.grad).to(self.device)
 
-        torch_tensor = torch.tensor([[[10, 10], [-4, -4]], [[5., 6], [7, 8]]], requires_grad=True).to(self.device)
+        torch_tensor = torch.tensor([[[10, 10], [-4, -4]], [[5., 6], [7, 8]]], requires_grad=True)
+
+        torch_tensor.to(self.device)
         
         torch_result = torch_tensor.min()
         torch_result.backward()
@@ -205,7 +223,9 @@ class TestTensorAutograd(unittest.TestCase):
         norch_result.backward()
         norch_tensor_grad = utils.to_torch(norch_tensor.grad).to(self.device)
 
-        torch_tensor = torch.tensor([[[10, 10], [-4, -4]], [[5., 6], [7, 8]]], requires_grad=True).to(self.device)
+        torch_tensor = torch.tensor([[[10, 10], [-4, -4]], [[5., 6], [7, 8]]], requires_grad=True)
+
+        torch_tensor.to(self.device)
         
         torch_min, _ = torch_tensor.min(axis=1)
         torch_result = torch_min.sum()
@@ -221,7 +241,9 @@ class TestTensorAutograd(unittest.TestCase):
         norch_result.backward()
         norch_tensor_grad = utils.to_torch(norch_tensor.grad).to(self.device)
 
-        torch_tensor = torch.tensor([[[10, 1], [-4, 0]], [[5., 50], [7, 8]]], requires_grad=True).to(self.device)
+        torch_tensor = torch.tensor([[[10, 1], [-4, 0]], [[5., 50], [7, 8]]], requires_grad=True)
+
+        torch_tensor.to(self.device)
         
         torch_min, _ = torch_tensor.min(axis=2)
         torch_result = torch_min.sum()
@@ -242,8 +264,12 @@ class TestTensorAutograd(unittest.TestCase):
         norch_tensor1_grad = utils.to_torch(norch_tensor1.grad).to(self.device)
         norch_tensor2_grad = utils.to_torch(norch_tensor2.grad).to(self.device)
 
-        torch_tensor1 = torch.tensor([[[1., 2, 3], [4, 5, 6]]], requires_grad=True).to(self.device)  # Shape (1, 2, 3)
-        torch_tensor2 = torch.tensor([1.5, -1, 0], requires_grad=True).to(self.device)  # Shape (3)
+        torch_tensor1 = torch.tensor([[[1., 2, 3], [4, 5, 6]]], requires_grad=True)  # Shape (1, 2, 3)
+        torch_tensor2 = torch.tensor([1.5, -1, 0], requires_grad=True)  # Shape (3)
+
+        torch_tensor1.to(self.device)
+        torch_tensor2.to(self.device)
+
         torch_result = (torch_tensor1 + torch_tensor2).sum()
         torch_result.backward()
         torch_tensor1_grad = torch_tensor1.grad
@@ -261,8 +287,11 @@ class TestTensorAutograd(unittest.TestCase):
         norch_tensor1_grad = utils.to_torch(norch_tensor1.grad).to(self.device)
         norch_tensor2_grad = utils.to_torch(norch_tensor2.grad).to(self.device)
 
-        torch_tensor1 = torch.tensor([[[1., 2, 3], [4, 5, 6]]], requires_grad=True).to(self.device)  # Shape (1, 2, 3)
-        torch_tensor2 = torch.tensor([1.5, -1, 0], requires_grad=True).to(self.device)  # Shape (3)
+        torch_tensor1 = torch.tensor([[[1., 2, 3], [4, 5, 6]]], requires_grad=True)  # Shape (1, 2, 3)
+        torch_tensor2 = torch.tensor([1.5, -1, 0], requires_grad=True)  # Shape (3)
+
+        torch_tensor1.to(self.device)
+        torch_tensor2.to(self.device)
 
         torch_result = (torch_tensor2 + torch_tensor1).sum()
         torch_result.backward()
@@ -283,8 +312,12 @@ class TestTensorAutograd(unittest.TestCase):
         norch_tensor1_grad_sub = utils.to_torch(norch_tensor1_sub.grad).to(self.device)
         norch_tensor2_grad_sub = utils.to_torch(norch_tensor2_sub.grad).to(self.device)
 
-        torch_tensor1_sub = torch.tensor([[[1, 2.5], [3, -4]], [[5, 6], [7, 8]]], requires_grad=True).to(self.device)
-        torch_tensor2_sub = torch.tensor([[[1, 1.], [1, 1.9]], [[1, 1], [1, 1]]], requires_grad=True).to(self.device)
+        torch_tensor1_sub = torch.tensor([[[1, 2.5], [3, -4]], [[5, 6], [7, 8]]], requires_grad=True)
+        torch_tensor2_sub = torch.tensor([[[1, 1.], [1, 1.9]], [[1, 1], [1, 1]]], requires_grad=True)
+
+        torch_tensor1_sub.to(self.device)
+        torch_tensor2_sub.to(self.device)
+
         torch_result_sub = (torch_tensor1_sub - torch_tensor2_sub).sum()
         torch_result_sub.backward()
         torch_tensor1_grad_sub = torch_tensor1_sub.grad
@@ -304,8 +337,12 @@ class TestTensorAutograd(unittest.TestCase):
         norch_tensor1_grad = utils.to_torch(norch_tensor1.grad).to(self.device)
         norch_tensor2_grad = utils.to_torch(norch_tensor2.grad).to(self.device)
 
-        torch_tensor1 = torch.tensor([[[1., 2, 3], [4, 5, 6]]], requires_grad=True).to(self.device)  # Shape (1, 2, 3)
-        torch_tensor2 = torch.tensor([1.5, -1, 0], requires_grad=True).to(self.device)  # Shape (3)
+        torch_tensor1 = torch.tensor([[[1., 2, 3], [4, 5, 6]]], requires_grad=True)  # Shape (1, 2, 3)
+        torch_tensor2 = torch.tensor([1.5, -1, 0], requires_grad=True)  # Shape (3)
+
+        torch_tensor1.to(self.device)
+        torch_tensor2.to(self.device)
+
         torch_result = (torch_tensor1 - torch_tensor2).sum()
         torch_result.backward()
         torch_tensor1_grad = torch_tensor1.grad
@@ -323,8 +360,11 @@ class TestTensorAutograd(unittest.TestCase):
         norch_tensor1_grad = utils.to_torch(norch_tensor1.grad).to(self.device)
         norch_tensor2_grad = utils.to_torch(norch_tensor2.grad).to(self.device)
 
-        torch_tensor1 = torch.tensor([[[1., 2, 3], [4, 5, 6]]], requires_grad=True).to(self.device)  # Shape (1, 2, 3)
-        torch_tensor2 = torch.tensor([1.5, -1, 0], requires_grad=True).to(self.device)  # Shape (3)
+        torch_tensor1 = torch.tensor([[[1., 2, 3], [4, 5, 6]]], requires_grad=True)  # Shape (1, 2, 3)
+        torch_tensor2 = torch.tensor([1.5, -1, 0], requires_grad=True)  # Shape (3)
+
+        torch_tensor1.to(self.device)
+        torch_tensor2.to(self.device)
         
         torch_result = (torch_tensor2 - torch_tensor1).sum()
         torch_result.backward()
@@ -346,8 +386,12 @@ class TestTensorAutograd(unittest.TestCase):
         norch_tensor1_grad_div = utils.to_torch(norch_tensor1_div.grad).to(self.device)
         norch_tensor2_grad_div = utils.to_torch(norch_tensor2_div.grad).to(self.device)
 
-        torch_tensor1_div = torch.tensor([[[2, 5.1], [6, -8]], [[10, 12], [14, 16]]], requires_grad=True).to(self.device)
-        torch_tensor2_div = torch.tensor([[[1, 1], [2, 2.2]], [[3, 3], [4, 4]]], requires_grad=True).to(self.device)
+        torch_tensor1_div = torch.tensor([[[2, 5.1], [6, -8]], [[10, 12], [14, 16]]], requires_grad=True)
+        torch_tensor2_div = torch.tensor([[[1, 1], [2, 2.2]], [[3, 3], [4, 4]]], requires_grad=True)
+
+        torch_tensor1_div.to(self.device)
+        torch_tensor2_div.to(self.device)
+
         torch_result_div = (torch_tensor1_div / torch_tensor2_div).sum()
         torch_result_div.backward()
         torch_tensor1_grad_div = torch_tensor1_div.grad
@@ -367,7 +411,10 @@ class TestTensorAutograd(unittest.TestCase):
         norch_result_div_scalar.backward()
         norch_tensor_grad_div_scalar = utils.to_torch(norch_tensor_div_scalar.grad).to(self.device)
 
-        torch_tensor_div_scalar = torch.tensor([[[2, 4.7], [6, 8]], [[10, 12], [14, 16]]], requires_grad=True).to(self.device)
+        torch_tensor_div_scalar = torch.tensor([[[2, 4.7], [6, 8]], [[10, 12], [14, 16]]], requires_grad=True)
+
+        torch_tensor_div_scalar.to(self.device)
+        
         torch_result_div_scalar = (torch_tensor_div_scalar / scalar).sum()
         torch_result_div_scalar.backward()
         torch_tensor_grad_div_scalar = torch_tensor_div_scalar.grad
@@ -385,7 +432,10 @@ class TestTensorAutograd(unittest.TestCase):
         norch_result_scalar_div.backward()
         norch_tensor_grad_scalar_div = utils.to_torch(norch_tensor_scalar_div.grad).to(self.device)
 
-        torch_tensor_scalar_div = torch.tensor([[[1, 2.23], [3, 4]], [[5, 6], [7, 8]]], requires_grad=True).to(self.device)
+        torch_tensor_scalar_div = torch.tensor([[[1, 2.23], [3, 4]], [[5, 6], [7, 8]]], requires_grad=True)
+
+        torch_tensor_scalar_div.to(self.device)
+        
         torch_result_scalar_div = (scalar / torch_tensor_scalar_div).sum()
         torch_result_scalar_div.backward()
         torch_tensor_grad_scalar_div = torch_tensor_scalar_div.grad
@@ -403,7 +453,10 @@ class TestTensorAutograd(unittest.TestCase):
         norch_result_power_st.backward()
         norch_tensor_grad_power_st = utils.to_torch(norch_tensor_power_st.grad).to(self.device)
 
-        torch_tensor_power_st = torch.tensor([[[2, 3.21], [4, 2.1]], [[6, 7], [8, 9]]], requires_grad=True).to(self.device)
+        torch_tensor_power_st = torch.tensor([[[2, 3.21], [4, 2.1]], [[6, 7], [8, 9]]], requires_grad=True)
+
+        torch_tensor_power_st.to(self.device)
+        
         torch_result_power_st = (scalar ** torch_tensor_power_st).sum()
         torch_result_power_st.backward()
         torch_tensor_grad_power_st = torch_tensor_power_st.grad
@@ -420,7 +473,10 @@ class TestTensorAutograd(unittest.TestCase):
         norch_result_power_ts.backward()
         norch_tensor_grad_power_ts = utils.to_torch(norch_tensor_power_ts.grad).to(self.device)
 
-        torch_tensor_power_ts = torch.tensor([[[2, 3], [4, 2.1]], [[6, 7], [8, 9]]], requires_grad=True).to(self.device)
+        torch_tensor_power_ts = torch.tensor([[[2, 3], [4, 2.1]], [[6, 7], [8, 9]]], requires_grad=True)
+
+        torch_tensor_power_ts.to(self.device)
+        
         torch_result_power_ts = (torch_tensor_power_ts ** scalar).sum()
         torch_result_power_ts.backward()
         torch_tensor_grad_power_ts = torch_tensor_power_ts.grad
@@ -438,8 +494,12 @@ class TestTensorAutograd(unittest.TestCase):
         norch_tensor1_grad_matmul = utils.to_torch(norch_tensor1_matmul.grad).to(self.device)
         norch_tensor2_grad_matmul = utils.to_torch(norch_tensor2_matmul.grad).to(self.device)
 
-        torch_tensor1_matmul = torch.tensor([[[1, 2.1], [3, -4]], [[5, 6], [7, 8]]], requires_grad=True).to(self.device)
-        torch_tensor2_matmul = torch.tensor([[[1.1, 3], [4, 5]], [[6, 7], [8, 9]]], requires_grad=True).to(self.device)
+        torch_tensor1_matmul = torch.tensor([[[1, 2.1], [3, -4]], [[5, 6], [7, 8]]], requires_grad=True)
+        torch_tensor2_matmul = torch.tensor([[[1.1, 3], [4, 5]], [[6, 7], [8, 9]]], requires_grad=True)
+
+        torch_tensor1_matmul.to(self.device)
+        torch_tensor2_matmul.to(self.device)
+
         torch_result_matmul = (torch_tensor1_matmul @ torch_tensor2_matmul).sum()
         torch_result_matmul.backward()
         torch_tensor1_grad_matmul = torch_tensor1_matmul.grad
@@ -466,8 +526,12 @@ class TestTensorAutograd(unittest.TestCase):
         norch_tensor2_grad_matmul = utils.to_torch(norch_tensor2_matmul.grad).to(self.device)
 
         # Repeat the same process with torch tensors
-        torch_tensor1_matmul = torch.tensor([[[1., 2], [3, -4], [5, 6], [7, 8]] for _ in range(B)], requires_grad=True).to(self.device)
-        torch_tensor2_matmul = torch.tensor([[[2., 3, 1, 0, 4], [5, -1, 2, 3, 0]] for _ in range(B)], requires_grad=True).to(self.device)
+        torch_tensor1_matmul = torch.tensor([[[1., 2], [3, -4], [5, 6], [7, 8]] for _ in range(B)], requires_grad=True)
+        torch_tensor2_matmul = torch.tensor([[[2., 3, 1, 0, 4], [5, -1, 2, 3, 0]] for _ in range(B)], requires_grad=True)
+
+        torch_tensor1_matmul.to(self.device)
+        torch_tensor2_matmul.to(self.device)
+
         torch_result_matmul = torch.matmul(torch_tensor1_matmul, torch_tensor2_matmul)
         torch_result_matmul_sum = torch_result_matmul.sum()
         torch_result_matmul_sum.backward()
@@ -498,8 +562,12 @@ class TestTensorAutograd(unittest.TestCase):
         norch_tensor2_grad_matmul = utils.to_torch(norch_tensor2_matmul.grad).to(self.device)
 
         # Repeat the same process with torch tensors
-        torch_tensor1_matmul = torch.tensor([[1., 2], [3, -4], [5, 6], [7, 8]], requires_grad=True).to(self.device)
-        torch_tensor2_matmul = torch.tensor([[[2., 3, 1, 0, 4], [5, -1, 2, 3, 0]] for _ in range(B)], requires_grad=True).to(self.device)
+        torch_tensor1_matmul = torch.tensor([[1., 2], [3, -4], [5, 6], [7, 8]], requires_grad=True)
+        torch_tensor2_matmul = torch.tensor([[[2., 3, 1, 0, 4], [5, -1, 2, 3, 0]] for _ in range(B)], requires_grad=True)
+
+        torch_tensor1_matmul.to(self.device)
+        torch_tensor2_matmul.to(self.device)
+
         torch_result_matmul = torch.matmul(torch_tensor1_matmul, torch_tensor2_matmul)
         torch_result_matmul_sum = torch_result_matmul.sum()
         torch_result_matmul_sum.backward()
@@ -523,7 +591,10 @@ class TestTensorAutograd(unittest.TestCase):
         norch_result_elemwise_mul_scalar.backward()
         norch_tensor_grad_elemwise_mul_scalar = utils.to_torch(norch_tensor_elemwise_mul_scalar.grad).to(self.device)
 
-        torch_tensor_elemwise_mul_scalar = torch.tensor([[[1.1, 2], [3, -4]], [[5, 6], [7, 8]]], requires_grad=True).to(self.device)
+        torch_tensor_elemwise_mul_scalar = torch.tensor([[[1.1, 2], [3, -4]], [[5, 6], [7, 8]]], requires_grad=True)
+
+        torch_tensor_elemwise_mul_scalar.to(self.device)
+        
         torch_result_elemwise_mul_scalar = (scalar * torch_tensor_elemwise_mul_scalar).sum()
         torch_result_elemwise_mul_scalar.backward()
         torch_tensor_grad_elemwise_mul_scalar = torch_tensor_elemwise_mul_scalar.grad
@@ -544,6 +615,10 @@ class TestTensorAutograd(unittest.TestCase):
 
         torch_tensor1_elemwise_mul = torch.tensor([[[1, 2.1], [3, -4]], [[5, 6], [7, 8]]], requires_grad=True).to(self.device)
         torch_tensor2_elemwise_mul = torch.tensor([[[1.1, 3], [4, 5]], [[6, 7], [8, 9]]], requires_grad=True).to(self.device)
+
+        torch_tensor1_elemwise_mul.to(self.device)
+        torch_tensor2_elemwise_mul.to(self.device)
+
         torch_result_elemwise_mul = (torch_tensor1_elemwise_mul * torch_tensor2_elemwise_mul).sum()
         torch_result_elemwise_mul.backward()
         torch_tensor1_grad_elemwise_mul = torch_tensor1_elemwise_mul.grad
@@ -561,7 +636,10 @@ class TestTensorAutograd(unittest.TestCase):
         norch_result_sin_tensor.backward()
         torch_result_sin_tensor_grad = utils.to_torch(norch_sin_tensor.grad).to(self.device)
 
-        torch_sin_tensor = torch.tensor([[[2, 3.21], [4, 2.1]], [[6, 7], [8, 9]]], requires_grad=True).to(self.device)
+        torch_sin_tensor = torch.tensor([[[2, 3.21], [4, 2.1]], [[6, 7], [8, 9]]], requires_grad=True)
+        
+        torch_sin_tensor.to(self.device)
+
         torch_expected_sin_tensor = (torch.sin(torch_sin_tensor)).sum()
         torch_expected_sin_tensor.backward()
         torch_expected_sin_tensor_grad = torch_sin_tensor.grad
@@ -577,7 +655,10 @@ class TestTensorAutograd(unittest.TestCase):
         norch_result_cos_tensor.backward()
         torch_result_cos_tensor_grad = utils.to_torch(norch_cos_tensor.grad).to(self.device)
 
-        torch_cos_tensor = torch.tensor([[[2, 3.21], [4, 2.1]], [[6, 7], [8, 9]]], requires_grad=True).to(self.device)
+        torch_cos_tensor = torch.tensor([[[2, 3.21], [4, 2.1]], [[6, 7], [8, 9]]], requires_grad=True)
+
+        torch_cos_tensor.to(self.device)
+        
         torch_expected_cos_tensor = (torch.sin(torch_cos_tensor)).sum()
         torch_expected_cos_tensor.backward()
         torch_expected_cos_tensor_grad = torch_cos_tensor.grad
@@ -595,7 +676,10 @@ class TestTensorAutograd(unittest.TestCase):
         norch_result.backward()
         norch_tensor_grad = utils.to_torch(norch_tensor.grad).to(self.device)
 
-        torch_tensor = torch.tensor([[[10, 10], [-4, -4]], [[5., 6], [7, 8]]], requires_grad=True).to(self.device)
+        torch_tensor = torch.tensor([[[10, 10], [-4, -4]], [[5., 6], [7, 8]]], requires_grad=True)
+
+        torch_tensor.to(self.device)
+        
         torch_sigmoid = torch.sigmoid(torch_tensor)
         torch_result = torch_sigmoid.sum()
 
@@ -617,8 +701,12 @@ class TestTensorAutograd(unittest.TestCase):
         loss_norch.backward()  # Backpropagate the loss
         grad_norch = predictions_norch.grad
 
-        predictions_torch = torch.tensor([1.1, 2, 3, 4], requires_grad=True).to(self.device)
-        labels_torch = torch.tensor([4, 3, 2.1, 1]).to(self.device)
+        predictions_torch = torch.tensor([1.1, 2, 3, 4], requires_grad=True)
+        labels_torch = torch.tensor([4, 3, 2.1, 1])
+
+        predictions_torch.to(self.device)
+        labels_torch.to(self.device)
+
         loss_torch_expected = loss_fn_torch(predictions_torch, labels_torch)
         loss_torch_expected.backward()  # Backpropagate the loss
         grad_torch_expected = predictions_torch.grad
@@ -643,8 +731,12 @@ class TestTensorAutograd(unittest.TestCase):
         loss_norch.backward()  # Backpropagate the loss
         grad_norch = predictions_norch.grad
 
-        predictions_torch = torch.tensor([2.0, 1.0, 0.1], requires_grad=True).to(self.device)
-        labels_torch = torch.tensor(0).to(self.device)
+        predictions_torch = torch.tensor([2.0, 1.0, 0.1], requires_grad=True)
+        labels_torch = torch.tensor(0)
+
+        predictions_torch.to(self.device)
+        labels_torch.to(self.device)
+
         loss_torch_expected = loss_fn_torch(predictions_torch, labels_torch)
         loss_torch_expected.backward()  # Backpropagate the loss
         grad_torch_expected = predictions_torch.grad
@@ -661,8 +753,12 @@ class TestTensorAutograd(unittest.TestCase):
         loss_norch.backward()  # Backpropagate the loss
         grad_norch = predictions_norch.grad
 
-        predictions_torch = torch.tensor([[0.5, 1.5, 2.5], [1.0, 2.0, 3.0]], requires_grad=True).to(self.device)
-        labels_torch = torch.tensor([2, 1]).to(self.device)
+        predictions_torch = torch.tensor([[0.5, 1.5, 2.5], [1.0, 2.0, 3.0]], requires_grad=True)
+        labels_torch = torch.tensor([2, 1])
+
+        predictions_torch.to(self.device)
+        labels_torch.to(self.device)
+
         loss_torch_expected = loss_fn_torch(predictions_torch, labels_torch)
         loss_torch_expected.backward()  # Backpropagate the loss
         grad_torch_expected = predictions_torch.grad
@@ -679,8 +775,12 @@ class TestTensorAutograd(unittest.TestCase):
         loss_norch.backward()  # Backpropagate the loss
         grad_norch = predictions_norch.grad
 
-        predictions_torch = torch.tensor([[0.0, 0.0, 0.0], [0.0, 0.0, 0.0]], requires_grad=True).to(self.device)
-        labels_torch = torch.tensor([1, 2]).to(self.device)
+        predictions_torch = torch.tensor([[0.0, 0.0, 0.0], [0.0, 0.0, 0.0]], requires_grad=True)
+        labels_torch = torch.tensor([1, 2])
+
+        predictions_torch.to(self.device)
+        labels_torch.to(self.device)
+
         loss_torch_expected = loss_fn_torch(predictions_torch, labels_torch)
         loss_torch_expected.backward()  # Backpropagate the loss
         grad_torch_expected = predictions_torch.grad
@@ -697,8 +797,12 @@ class TestTensorAutograd(unittest.TestCase):
         loss_norch.backward()  # Backpropagate the loss
         grad_norch = predictions_norch.grad
 
-        predictions_torch = torch.tensor([[0.5, 0.2, 0.1], [0.1, 0.5, 0.7]], requires_grad=True).to(self.device)
-        labels_torch = torch.tensor([[1., 0, 0], [0, 1, 0]]).to(self.device)
+        predictions_torch = torch.tensor([[0.5, 0.2, 0.1], [0.1, 0.5, 0.7]], requires_grad=True)
+        labels_torch = torch.tensor([[1., 0, 0], [0, 1, 0]])
+
+        predictions_torch.to(self.device)
+        labels_torch.to(self.device)
+
         loss_torch_expected = loss_fn_torch(predictions_torch, labels_torch)
         loss_torch_expected.backward()  # Backpropagate the loss
         grad_torch_expected = predictions_torch.grad
@@ -757,7 +861,10 @@ class TestTensorAutograd(unittest.TestCase):
         norch_result_reshape.backward()
         norch_tensor_grad_reshape = utils.to_torch(norch_tensor_reshape.grad).to(self.device)
 
-        torch_tensor_reshape = torch.tensor([[[1, 2.1], [3, -4]], [[5, 6], [7, 8]]], requires_grad=True).to(self.device)
+        torch_tensor_reshape = torch.tensor([[[1, 2.1], [3, -4]], [[5, 6], [7, 8]]], requires_grad=True)
+
+        torch_tensor_reshape.to(self.device)
+        
         torch_result_reshape = torch_tensor_reshape.reshape(new_shape).sum()
         torch_result_reshape.backward()
         torch_tensor_grad_reshape = torch_tensor_reshape.grad
@@ -775,7 +882,10 @@ class TestTensorAutograd(unittest.TestCase):
         norch_result_transpose.backward()
         norch_tensor_grad_transpose = utils.to_torch(norch_tensor_transpose.grad).to(self.device)
 
-        torch_tensor_transpose = torch.tensor([[[1, 2.1], [3, -4]], [[5, 6], [7, 8]]], requires_grad=True).to(self.device)
+        torch_tensor_transpose = torch.tensor([[[1, 2.1], [3, -4]], [[5, 6], [7, 8]]], requires_grad=True)
+
+        torch_tensor_transpose.to(self.device)
+
         torch_result_transpose = torch_tensor_transpose.transpose(axis1, axis2).sum()
         torch_result_transpose.backward()
         torch_tensor_grad_transpose = torch_tensor_transpose.grad
@@ -792,7 +902,10 @@ class TestTensorAutograd(unittest.TestCase):
         norch_result_T.backward()
         norch_tensor_grad_T = utils.to_torch(norch_tensor_T.grad).to(self.device)
 
-        torch_tensor_T = torch.tensor([[[1, 2.1], [3, -4]], [[5, 6], [7, 8]]], requires_grad=True).to(self.device)
+        torch_tensor_T = torch.tensor([[[1, 2.1], [3, -4]], [[5, 6], [7, 8]]], requires_grad=True)
+
+        torch_tensor_T.to(self.device)
+        
         torch_result_T = torch_tensor_T.mT.sum()
         torch_result_T.backward()
         torch_tensor_grad_T = torch_tensor_T.grad
@@ -813,8 +926,11 @@ class TestTensorAutograd(unittest.TestCase):
         norch_tensor_grad_reshape_matmul1 = utils.to_torch(norch_tensor1.grad).to(self.device)
         norch_tensor_grad_reshape_matmul2 = utils.to_torch(norch_tensor2.grad).to(self.device)
 
-        torch_tensor1 = torch.tensor([[1, 2.1], [3, -4], [5, 6], [7, 8]], dtype=torch.float32, requires_grad=True).to(self.device)
-        torch_tensor2 = torch.tensor([[1, 5.1], [0.1, -4], [0, 6], [7, 8]], dtype=torch.float32, requires_grad=True).to(self.device)     
+        torch_tensor1 = torch.tensor([[1, 2.1], [3, -4], [5, 6], [7, 8]], dtype=torch.float32, requires_grad=True)
+        torch_tensor2 = torch.tensor([[1, 5.1], [0.1, -4], [0, 6], [7, 8]], dtype=torch.float32, requires_grad=True)
+
+        torch_tensor1.to(self.device)
+        torch_tensor2.to(self.device)
         
         torch_result_reshape_matmul = (torch_tensor1.reshape(new_shape) @ torch_tensor2).sum()
         torch_result_reshape_matmul.backward()
@@ -835,7 +951,10 @@ class TestTensorAutograd(unittest.TestCase):
         norch_result_unsqueeze_0.backward()
         norch_tensor_grad_unsqueeze_0 = utils.to_torch(norch_tensor_unsqueeze.grad).to(self.device)
 
-        torch_tensor_unsqueeze = torch.tensor([[1., 2], [3, 4]], requires_grad=True).to(self.device)
+        torch_tensor_unsqueeze = torch.tensor([[1., 2], [3, 4]], requires_grad=True)
+
+        torch_tensor_unsqueeze.to(self.device)
+        
         torch_result_unsqueeze_0 = torch_tensor_unsqueeze.unsqueeze(0).sum()
         torch_result_unsqueeze_0.backward()
         torch_tensor_grad_unsqueeze_0 = torch_tensor_unsqueeze.grad
@@ -848,7 +967,10 @@ class TestTensorAutograd(unittest.TestCase):
         norch_result_unsqueeze_1.backward()
         norch_tensor_grad_unsqueeze_1 = utils.to_torch(norch_tensor_unsqueeze.grad).to(self.device)
 
-        torch_tensor_unsqueeze = torch.tensor([[1., 2.], [3, 4]], requires_grad=True).to(self.device)
+        torch_tensor_unsqueeze = torch.tensor([[1., 2.], [3, 4]], requires_grad=True)
+        
+        torch_tensor_unsqueeze.to(self.device)
+
         torch_result_unsqueeze_1 = torch_tensor_unsqueeze.unsqueeze(1).sum()
         torch_result_unsqueeze_1.backward()
         torch_tensor_grad_unsqueeze_1 = torch_tensor_unsqueeze.grad
@@ -861,7 +983,10 @@ class TestTensorAutograd(unittest.TestCase):
         norch_result_unsqueeze_2.backward()
         norch_tensor_grad_unsqueeze_2 = utils.to_torch(norch_tensor_unsqueeze.grad).to(self.device)
 
-        torch_tensor_unsqueeze = torch.tensor([[1., 2], [3, 4]], requires_grad=True).to(self.device)
+        torch_tensor_unsqueeze = torch.tensor([[1., 2], [3, 4]], requires_grad=True)
+        
+        torch_tensor_unsqueeze.to(self.device)
+
         torch_result_unsqueeze_2 = torch_tensor_unsqueeze.unsqueeze(2).sum()
         torch_result_unsqueeze_2.backward()
         torch_tensor_grad_unsqueeze_2 = torch_tensor_unsqueeze.grad
@@ -881,8 +1006,11 @@ class TestTensorAutograd(unittest.TestCase):
         norch_tensor_grad_unsqueeze_matmul1 = utils.to_torch(norch_tensor1.grad).to(self.device)
         norch_tensor_grad_unsqueeze_matmul2 = utils.to_torch(norch_tensor2.grad).to(self.device)
 
-        torch_tensor1 = torch.tensor([[1, 2], [3, 4]], dtype=torch.float32, requires_grad=True).to(self.device)
-        torch_tensor2 = torch.tensor([[1, 2], [3, 4]], dtype=torch.float32, requires_grad=True).to(self.device)
+        torch_tensor1 = torch.tensor([[1, 2], [3, 4]], dtype=torch.float32, requires_grad=True)
+        torch_tensor2 = torch.tensor([[1, 2], [3, 4]], dtype=torch.float32, requires_grad=True)
+
+        torch_tensor1.to(self.device)
+        torch_tensor2.to(self.device)
         
         torch_result_unsqueeze_matmul = (torch_tensor1 @ torch_tensor2.unsqueeze(0)).sum()
         torch_result_unsqueeze_matmul.backward()
@@ -902,7 +1030,10 @@ class TestTensorAutograd(unittest.TestCase):
         norch_result_squeeze_0.backward()
         norch_tensor_grad_squeeze_0 = utils.to_torch(norch_tensor_squeeze.grad).to(self.device)
 
-        torch_tensor_squeeze = torch.tensor([[[1., 2], [3, 4]]], requires_grad=True).to(self.device)
+        torch_tensor_squeeze = torch.tensor([[[1., 2], [3, 4]]], requires_grad=True)
+
+        torch_tensor_squeeze.to(self.device)
+        
         torch_result_squeeze_0 = torch_tensor_squeeze.squeeze(0).sum()
         torch_result_squeeze_0.backward()
         torch_tensor_grad_squeeze_0 = torch_tensor_squeeze.grad
@@ -922,8 +1053,11 @@ class TestTensorAutograd(unittest.TestCase):
         norch_tensor_grad_squeeze_matmul1 = utils.to_torch(norch_tensor1.grad).to(self.device)
         norch_tensor_grad_squeeze_matmul2 = utils.to_torch(norch_tensor2.grad).to(self.device)
 
-        torch_tensor1 = torch.tensor([[[1., 2], [3, 4]]], dtype=torch.float32, requires_grad=True).to(self.device)
-        torch_tensor2 = torch.tensor([[[1., 2], [3, 4]]], dtype=torch.float32, requires_grad=True).to(self.device)
+        torch_tensor1 = torch.tensor([[[1., 2], [3, 4]]], dtype=torch.float32, requires_grad=True)
+        torch_tensor2 = torch.tensor([[[1., 2], [3, 4]]], dtype=torch.float32, requires_grad=True)
+
+        torch_tensor1.to(self.device)
+        torch_tensor2.to(self.device)
         
         torch_result_squeeze_matmul = (torch_tensor1.squeeze(0) @ torch_tensor2).sum()
         torch_result_squeeze_matmul.backward()
@@ -946,8 +1080,11 @@ class TestTensorAutograd(unittest.TestCase):
         norch_tensor_grad_T_matmul1 = utils.to_torch(norch_tensor1.grad).to(self.device)
         norch_tensor_grad_T_matmult2 = utils.to_torch(norch_tensor2.grad).to(self.device)
 
-        torch_tensor1 = torch.tensor([[1, 2.1], [3, -4], [5, 6], [7, 8]], dtype=torch.float32, requires_grad=True).to(self.device)
-        torch_tensor2 = torch.tensor([[1, 5.1], [0.1, -4], [0, 6], [7, 8]], dtype=torch.float32, requires_grad=True).to(self.device)     
+        torch_tensor1 = torch.tensor([[1, 2.1], [3, -4], [5, 6], [7, 8]], dtype=torch.float32, requires_grad=True)
+        torch_tensor2 = torch.tensor([[1, 5.1], [0.1, -4], [0, 6], [7, 8]], dtype=torch.float32, requires_grad=True)
+
+        torch_tensor1.to(self.device)
+        torch_tensor2.to(self.device)
         
         torch_result_T_matmul = (torch_tensor1.T @ torch_tensor2).sum()
         torch_result_T_matmul.backward()
@@ -962,7 +1099,7 @@ class TestTensorAutograd(unittest.TestCase):
         The code has a problem on the following operation
         tensor1.reshape(..) @ tensor1
         print(tensor1.grad)
-        (also transpsoe and .T)
+        (also transpose and .T)
         """
         pass
 
@@ -978,8 +1115,11 @@ class TestTensorAutograd(unittest.TestCase):
         norch_tensor_grad_transpose_matmul1 = utils.to_torch(norch_tensor1.grad).to(self.device)
         norch_tensor_grad_transpose_matmult2 = utils.to_torch(norch_tensor2.grad).to(self.device)
 
-        torch_tensor1 = torch.tensor([[1, 2.1], [3, -4], [5, 6], [7, 8]], dtype=torch.float32, requires_grad=True).to(self.device)
-        torch_tensor2 = torch.tensor([[1, 5.1], [0.1, -4], [0, 6], [7, 8]], dtype=torch.float32, requires_grad=True).to(self.device)  
+        torch_tensor1 = torch.tensor([[1, 2.1], [3, -4], [5, 6], [7, 8]], dtype=torch.float32, requires_grad=True)
+        torch_tensor2 = torch.tensor([[1, 5.1], [0.1, -4], [0, 6], [7, 8]], dtype=torch.float32, requires_grad=True)
+
+        torch_tensor1.to(self.device)
+        torch_tensor2.to(self.device)
         
         torch_result_transpose_matmul = (torch_tensor1.T @ torch_tensor2).sum()
         torch_result_transpose_matmul.backward()
diff --git a/tests/test_nn.py b/tests/test_nn.py
index cd623f3..4896506 100644
--- a/tests/test_nn.py
+++ b/tests/test_nn.py
@@ -26,8 +26,12 @@ class TestNNModuleLoss(unittest.TestCase):
         loss_norch = loss_fn_norch.forward(predictions_norch, labels_norch)
         loss_torch_result = utils.to_torch(loss_norch).to(self.device)
 
-        predictions_torch = torch.tensor([[1.1, 2, 3, 4], [1.1, 2, 3, 4]]).to(self.device)
-        labels_torch = torch.tensor([[1.1, 2, 3, 4], [1.1, 2, 3, 3]]).to(self.device)
+        predictions_torch = torch.tensor([[1.1, 2, 3, 4], [1.1, 2, 3, 4]])
+        labels_torch = torch.tensor([[1.1, 2, 3, 4], [1.1, 2, 3, 3]])
+
+        predictions_torch.to(self.device)
+        labels_torch.to(self.device)
+
         loss_torch_expected = loss_fn_torch(predictions_torch, labels_torch)        
         
         self.assertTrue(utils.compare_torch(loss_torch_result, loss_torch_expected))
@@ -38,8 +42,12 @@ class TestNNModuleLoss(unittest.TestCase):
         loss_norch = loss_fn_norch.forward(predictions_norch, labels_norch)
         loss_torch_result = utils.to_torch(loss_norch).to(self.device)
 
-        predictions_torch = torch.tensor([1.1, 2, 3, 4]).to(self.device)
-        labels_torch = torch.tensor([4, 3, 2.1, 1]).to(self.device)
+        predictions_torch = torch.tensor([1.1, 2, 3, 4])
+        labels_torch = torch.tensor([4, 3, 2.1, 1])
+
+        predictions_torch.to(self.device)
+        labels_torch.to(self.device)
+        
         loss_torch_expected = loss_fn_torch(predictions_torch, labels_torch)        
         
         self.assertTrue(utils.compare_torch(loss_torch_result, loss_torch_expected))
@@ -58,8 +66,12 @@ class TestNNModuleLoss(unittest.TestCase):
 
         loss_torch_result = utils.to_torch(loss_norch).to(self.device)
 
-        predictions_torch = torch.tensor([2.0, 1.0, 0.1]).to(self.device)
-        labels_torch = torch.tensor(0).to(self.device)
+        predictions_torch = torch.tensor([2.0, 1.0, 0.1])
+        labels_torch = torch.tensor(0)
+
+        predictions_torch.to(self.device)
+        labels_torch.to(self.device)
+
         loss_torch_expected = loss_fn_torch(predictions_torch, labels_torch)
 
         self.assertTrue(utils.compare_torch(loss_torch_result, loss_torch_expected))
@@ -71,8 +83,12 @@ class TestNNModuleLoss(unittest.TestCase):
         loss_torch_result = utils.to_torch(loss_norch).to(self.device)
         
 
-        predictions_torch = torch.tensor([[0.5, 1.5, 2.5], [1.0, 2.0, 3.0]]).to(self.device)
-        labels_torch = torch.tensor([2, 1]).to(self.device)
+        predictions_torch = torch.tensor([[0.5, 1.5, 2.5], [1.0, 2.0, 3.0]])
+        labels_torch = torch.tensor([2, 1])
+
+        predictions_torch.to(self.device)
+        labels_torch.to(self.device)
+
         loss_torch_expected = loss_fn_torch(predictions_torch, labels_torch)
         
         self.assertTrue(utils.compare_torch(loss_torch_result, loss_torch_expected))
@@ -83,8 +99,12 @@ class TestNNModuleLoss(unittest.TestCase):
         loss_norch = loss_fn_norch.forward(predictions_norch, labels_norch)
         loss_torch_result = utils.to_torch(loss_norch).to(self.device)
         
-        predictions_torch = torch.tensor([[0.0, 0.0, 0.0], [0.0, 0.0, 0.0]]).to(self.device)
-        labels_torch = torch.tensor([1, 2]).to(self.device)
+        predictions_torch = torch.tensor([[0.0, 0.0, 0.0], [0.0, 0.0, 0.0]])
+        labels_torch = torch.tensor([1, 2])
+
+        predictions_torch.to(self.device)
+        labels_torch.to(self.device)
+
         loss_torch_expected = loss_fn_torch(predictions_torch, labels_torch)
         
         self.assertTrue(utils.compare_torch(loss_torch_result, loss_torch_expected))
@@ -95,8 +115,12 @@ class TestNNModuleLoss(unittest.TestCase):
         loss_norch = loss_fn_norch.forward(predictions_norch, labels_norch)
         loss_torch_result = utils.to_torch(loss_norch).to(self.device)
         
-        predictions_torch = torch.tensor([0.5, 0.2, 0.1]).to(self.device)
-        labels_torch = torch.tensor([1., 0, 0]).to(self.device)
+        predictions_torch = torch.tensor([0.5, 0.2, 0.1])
+        labels_torch = torch.tensor([1., 0, 0])
+
+        predictions_torch.to(self.device)
+        labels_torch.to(self.device)
+
         loss_torch_expected = loss_fn_torch(predictions_torch, labels_torch)
         
         self.assertTrue(utils.compare_torch(loss_torch_result, loss_torch_expected))
@@ -107,8 +131,12 @@ class TestNNModuleLoss(unittest.TestCase):
         loss_norch = loss_fn_norch.forward(predictions_norch, labels_norch)
         loss_torch_result = utils.to_torch(loss_norch).to(self.device)
 
-        predictions_torch = torch.tensor([[0.5, 0.2, 0.1], [0.1, 0.5, 0.7]]).to(self.device)
-        labels_torch = torch.tensor([[1., 0, 0], [0, 1, 0]]).to(self.device)
+        predictions_torch = torch.tensor([[0.5, 0.2, 0.1], [0.1, 0.5, 0.7]])
+        labels_torch = torch.tensor([[1., 0, 0], [0, 1, 0]])
+
+        predictions_torch.to(self.device)
+        labels_torch.to(self.device)
+
         loss_torch_expected = loss_fn_torch(predictions_torch, labels_torch)
         
         self.assertTrue(utils.compare_torch(loss_torch_result, loss_torch_expected))
@@ -133,7 +161,10 @@ class TestNNModuleActivationFn(unittest.TestCase):
         sigmoid_norch = sigmoid_fn_norch.forward(x)
         sigmoid_torch_result = utils.to_torch(sigmoid_norch).to(self.device)
 
-        x = torch.tensor([[1, 2, 3]]).to(self.device)
+        x = torch.tensor([[1, 2, 3]])
+
+        x.to(self.device)
+        
         sigmoid_torch_expected = sigmoid_fn_torch.forward(x)
 
         self.assertTrue(utils.compare_torch(sigmoid_torch_result, sigmoid_torch_expected))
@@ -143,7 +174,10 @@ class TestNNModuleActivationFn(unittest.TestCase):
         sigmoid_norch = sigmoid_fn_norch.forward(x)
         sigmoid_torch_result = utils.to_torch(sigmoid_norch).to(self.device)
 
-        x = torch.tensor([-1, 2, -3]).to(self.device)
+        x = torch.tensor([-1, 2, -3])
+
+        x.to(self.device)
+        
         sigmoid_torch_expected = sigmoid_fn_torch.forward(x)
 
         self.assertTrue(utils.compare_torch(sigmoid_torch_result, sigmoid_torch_expected))
@@ -153,7 +187,10 @@ class TestNNModuleActivationFn(unittest.TestCase):
         sigmoid_norch = sigmoid_fn_norch.forward(x)
         sigmoid_torch_result = utils.to_torch(sigmoid_norch).to(self.device)
 
-        x = torch.tensor([0, 0, 0]).to(self.device)
+        x = torch.tensor([0, 0, 0])
+
+        x.to(self.device)
+
         sigmoid_torch_expected = sigmoid_fn_torch.forward(x)
 
         self.assertTrue(utils.compare_torch(sigmoid_torch_result, sigmoid_torch_expected))
@@ -180,7 +217,9 @@ class TestNNModuleActivationFn(unittest.TestCase):
             for norch_input, torch_input in test_cases:
                 # Move tensors to the correct device
                 norch_input = norch_input.to(self.device)
-                torch_input = torch_input.to(self.device)
+                torch_input = torch_input
+                
+                torch_input.to(self.device)
 
                 # Forward pass using norch
                 softmax_norch = softmax_fn_norch.forward(norch_input)
diff --git a/tests/test_operations.py b/tests/test_operations.py
index d443260..262cee4 100644
--- a/tests/test_operations.py
+++ b/tests/test_operations.py
@@ -471,8 +471,33 @@ class TestTensorOperations(unittest.TestCase):
         
         self.assertTrue(utils.compare_torch(torch_result, torch_expected))
 
+    def test_1D_T(self):
+        """
+        Test transposition of a 1D tensor.
+        """
+        norch_tensor = norch.Tensor([1, 2, 3, 4]).to(self.device)
+        norch_result = norch_tensor.T
+        torch_result = utils.to_torch(norch_result).to(self.device)
 
-    def test_transpose_T(self):
+        torch_tensor = torch.tensor([1, 2, 3, 4]).to(self.device)
+        torch_expected = torch_tensor.T
+
+        self.assertTrue(utils.compare_torch(torch_result, torch_expected))
+
+    def test_2D_T(self):
+        """
+        Test transposition of a 2D tensor.
+        """
+        norch_tensor = norch.Tensor([[1, 2, 3], [4, 5, 6]]).to(self.device)
+        norch_result = norch_tensor.T
+        torch_result = utils.to_torch(norch_result).to(self.device)
+
+        torch_tensor = torch.tensor([[1, 2, 3], [4, 5, 6]]).to(self.device)
+        torch_expected = torch_tensor.T
+
+        self.assertTrue(utils.compare_torch(torch_result, torch_expected))
+
+    def test_3D_T(self):
         """
         Test transposition of a tensor: tensor.T
         """
@@ -481,7 +506,7 @@ class TestTensorOperations(unittest.TestCase):
         torch_result = utils.to_torch(norch_result).to(self.device)
 
         torch_tensor = torch.tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]]).to(self.device)
-        torch_expected = torch.transpose(torch_tensor, 0, 2)
+        torch_expected = torch_tensor.T
 
         self.assertTrue(utils.compare_torch(torch_result, torch_expected))