diff --git a/.gitignore b/.gitignore
deleted file mode 100644
index 2744889..0000000
--- a/.gitignore
+++ /dev/null
@@ -1,9 +0,0 @@
-build
-.vscode
-*.txt
-python_lib/dl_lib/_compiled
-*__pycache__*
-*_cache
-
-# TODO: remove later
-benchmarks
\ No newline at end of file
diff --git a/CMakeLists.txt b/CMakeLists.txt
index 90dba31..2936cac 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -27,7 +27,22 @@ endif()
 add_compile_options("$<$<C_COMPILER_ID:MSVC>:/utf-8>")
 add_compile_options("$<$<CXX_COMPILER_ID:MSVC>:/utf-8>")
 
-# TODO: add flag for double precision?
+option(CUDA "Enable CUDA execution for some faster data structures" ON)
+
+if (CUDA)  
+  include(CheckLanguage)
+  check_language(CUDA)
+
+  if(CMAKE_CUDA_COMPILER)
+    add_definitions(-D__CUDA)
+    enable_language(CUDA)
+    
+    set(CMAKE_CUDA_STANDARD 20)
+    set(CMAKE_CUDA_STANDARD_REQUIRED ON)
+  else()
+    message(WARNING "Could not find CUDA on system. Compiling without CUDA enabled")
+  endif()
+endif()
 
 # include python libs
 if(APPLE)
diff --git a/examples/mnist.py b/examples/mnist.py
index b0368e8..5a3ca9a 100644
--- a/examples/mnist.py
+++ b/examples/mnist.py
@@ -125,10 +125,10 @@ def evaluate(net, x, y_int, batch_size=256):
     # setup
     net = make_net()
     loss_fn = CrossEntropyWithSoftmax()
-    optim = RmsProp(net.parameters(), 0.00001, 0.95)  # lr and decay
+    optim = RmsProp(net.parameters(), 0.000001, 0.999)  # lr and decay
 
     # training loop
-    n_epochs = 10
+    n_epochs = 5
     for epoch in range(n_epochs):
         train_loss = train_epoch(net, loss_fn, optim, x_train, y_train)
         val_acc = evaluate(net, x_val, y_val)
diff --git a/readme.md b/readme.md
index 2a8152e..b7b137e 100644
--- a/readme.md
+++ b/readme.md
@@ -20,17 +20,18 @@ For some examples on Python interface, see tests/python.
   - Training framework (optimizers, loss functions, layers, and networks)
 - **Example code**: Full MNIST dataset training example
 - **Python Interface**: Seamless integration via Boost.Python
-- **Clean Architecture**: Modular design, ~4K LOC
+- **Clean Architecture**: Modular design, maintainable and extensible
 - **CI/CD**: Automated testing with GTest and GitHub Actions
 
 ## Tech Stack
 
-- C++17/20
+- C++17/20/23
 - CMake build system
 - Boost.Python for Python bindings
 - Python 3 for library interface and examples
 - Google Test (GTest) and PyTest for unit testing
 - GitHub Actions for CI/CD
+- CUDA
 
 ## Current Status
 
@@ -40,7 +41,7 @@ Roadmap:
 - [x] Python Binding Unit Tests
 - [x] Optimizers and training framework
 - [x] MNIST example
-- [ ] CUDA mode for operations
+- [x] CUDA mode for operations
 - [ ] Additional layer types (Conv2D, Dropout, etc.)
 - [ ] AlexNet reference implementation
 - [ ] Docker deployment example
@@ -50,9 +51,18 @@ Roadmap:
 mkdir build && cd build
 cmake ..
 make
-ctest
 ```
 
+### Building with CUDA
+
+Project automatically detects whether CUDA is installed, and compiles with it.
+If CUDA compilation not desired you can switch it off via 
+
+```bash
+cmake --DCUDA=Off ..
+```
+
+
 ## Running Unit Tests
 
 Compile with building tests enabled: 
@@ -61,7 +71,7 @@ Compile with building tests enabled:
 mkdir build && cd build
 cmake -DBUILD_TESTS=On ..
 make
-ctest 
+ctest .
 ```
 
 ## Required
@@ -73,6 +83,7 @@ ctest
 - numpy 1.26.4
 - pytest and GTest for unit tests (we use pytest=9.0.2)
 - Google Benchmark for benchmarking
+- CUDA (we use CUDA 13.1 on an RTX-5050)
 
 ## Troubleshooting
 
diff --git a/src/backend/CMakeLists.txt b/src/backend/CMakeLists.txt
index ed6bade..a4f89ed 100644
--- a/src/backend/CMakeLists.txt
+++ b/src/backend/CMakeLists.txt
@@ -7,6 +7,18 @@ file(GLOB_RECURSE CORE_SOURCES
     utility/*.cpp
 )
 
+if(CMAKE_CUDA_COMPILER)
+    file(GLOB_RECURSE CUDA_SOURCES
+        computational_graph/*.cu
+        data_modeling/*.cu
+        module/*.cu
+        system/*.cu
+        training/*.cu
+        utility/*.cu
+    )
+    list(APPEND CORE_SOURCES ${CUDA_SOURCES})
+endif()
+
 add_library(BackendCore SHARED ${CORE_SOURCES})
 
 target_include_directories(BackendCore PUBLIC
@@ -15,4 +27,23 @@ target_include_directories(BackendCore PUBLIC
 
 set_target_properties(BackendCore PROPERTIES
     LIBRARY_OUTPUT_DIRECTORY "${PYTHON_MODULE_DIR}" # make sure Python-modules see backend
-)
\ No newline at end of file
+)
+
+if(CMAKE_CUDA_COMPILER)
+    set_target_properties(BackendCore PROPERTIES
+        CUDA_SEPARABLE_COMPILATION ON
+        # nvidia-smi --query-gpu=compute_cap --format=csv,noheader 
+        # I get 12.0, hence 120
+        #set(CMAKE_CUDA_ARCHITECTURES "75;86;89;100;120")
+        CMAKE_CUDA_ARCHITECTURES native
+    )
+
+    find_package(CUDAToolkit REQUIRED)
+    target_include_directories(BackendCore PRIVATE 
+        ${CUDAToolkit_INCLUDE_DIRS}
+    )
+    target_link_libraries(BackendCore 
+        CUDA::cudart
+    )
+endif()
+
diff --git a/src/backend/computational_graph/activation_functions/cuda/activation_nodes.cu b/src/backend/computational_graph/activation_functions/cuda/activation_nodes.cu
new file mode 100644
index 0000000..e19165c
--- /dev/null
+++ b/src/backend/computational_graph/activation_functions/cuda/activation_nodes.cu
@@ -0,0 +1,214 @@
+/**
+ * @file activation_nodes.cu
+ * @author Robert Baumgartner (r.baumgartner-1@tudelft.nl)
+ * @brief
+ * @version 0.1
+ * @date 2026-03-23
+ *
+ * @copyright Copyright (c) 2026
+ *
+ */
+
+#ifndef __CUDA
+static_assert(false, "File should not be compiled without CUDA enabled");
+#endif // __CUDA
+
+#include "activation_nodes.cuh"
+#include "utility/cuda/cuda_common.cuh"
+
+using namespace std;
+
+namespace {
+  /**
+   * @brief Relu backward kernel.
+   */
+  __global__ void reluBackwardKernel(ftype* const res, const ftype* const upstreamGrad, const ftype* const parent, const tensorSize_t size) {
+    const int gid = blockIdx.x * blockDim.x + threadIdx.x;
+    if(gid >= size) {
+      return;
+    }
+
+    res[gid] =  parent[gid] > 0 ? upstreamGrad[gid] : 0;
+  }
+
+  /**
+   * @brief Leaky relu backward kernel.
+   */
+  __global__ void leakyReluBackwardKernel(ftype* const res, const ftype* const upstreamGrad, const ftype* const parent, const ftype eps, const tensorSize_t size) {
+    const int gid = blockIdx.x * blockDim.x + threadIdx.x;
+    if(gid >= size) {
+      return;
+    }
+
+    res[gid] = parent[gid] > 0 ? upstreamGrad[gid] : eps * upstreamGrad[gid];
+  }
+
+  /**
+   * @brief Sigmoid backward kernel, optimized by using the forward sigmoid.
+   */
+  __global__ void sigmoidBackwardKernel(ftype* const res, const ftype* const upstreamGrad, const ftype* const sigmoid, const tensorSize_t size) {
+    const int gid = blockIdx.x * blockDim.x + threadIdx.x;
+    if(gid >= size) {
+      return;
+    }
+
+    ftype si = sigmoid[gid];
+    res[gid] = si * (1 - si) * upstreamGrad[gid];
+  }
+
+  /**
+   * @brief Softmax backward kernel. This kernel is different than others since it is warp aligned. The inner loop avoids shared memory bank 
+   * conflicts by broadcasting.
+   * 
+   * stridesWidthPerBlock is an awkward name. It is the product of number of strides per block (times) stride. We pre-compute it on host. 
+   */
+  __global__ void softmaxBackwardKernelOneBlock(ftype* const res, const ftype* const upstreamGrad, const ftype* const softmax,  
+                                                const tensorSize_t stride, const int stridesWidthPerBlock, const int threadsPerStride, tensorSize_t size) {
+    const int tid = threadIdx.x;
+
+    const int withinStrideOffset = tid % threadsPerStride;
+    const int strideOffset = (tid / threadsPerStride) * stride;
+
+    const int gid = blockIdx.x * stridesWidthPerBlock + strideOffset + withinStrideOffset;
+    const bool isPadded = (withinStrideOffset >= stride) || (gid >= size); // padded threads only exists to align warps with strides
+
+    ftype yi = 0;
+    const int smemOffset = strideOffset + withinStrideOffset;
+
+    extern __shared__ ftype smem[];
+    if(!isPadded) {
+      yi = softmax[gid];
+      smem[smemOffset] = yi;
+      smem[smemOffset + stridesWidthPerBlock] = upstreamGrad[gid];
+    }
+    __syncthreads();
+
+    if(isPadded) {
+      return;
+    }
+
+    ftype grad = 0;
+    for(int j = 0; j < stride; j++) {
+      // warp alignment -> smem-reads are broadcasted per warp -> no bank conflicts
+      ftype yj = smem[strideOffset + j];
+      ftype gj = smem[strideOffset + j + stridesWidthPerBlock];
+
+      auto jacobian = (withinStrideOffset == j) ? yi * (1 - yj) : -yi * yj;
+      grad += gj * jacobian;
+    }
+
+    res[gid] = grad;
+  }
+
+  /**
+   * @brief Large softmax pass. Because the stride now does not fit into one block anymore we do a grid-stride loop.
+   */
+  __global__ void softmaxBackwardKernelLargePass(ftype* const res, const ftype* const upstreamGrad, const ftype* const softmax, const int blocksPerStride, const tensorSize_t stride) {    
+    const int strideNumber = blockIdx.x / blocksPerStride;
+    const int strideOffset = strideNumber * stride;
+    const int i = (blockIdx.x % blocksPerStride) * blockDim.x + threadIdx.x;
+    // blockIdx.x % blocksPerStride = block number within this stride
+
+    const int tid = threadIdx.x;
+    const int gid = strideOffset + i;
+
+    extern __shared__ ftype smem[];
+
+    const bool isNotPadded = i < stride;
+    const ftype yi = isNotPadded ? softmax[gid] : 0;
+
+    ftype grad = 0;
+    for(int offset = 0; offset < stride; offset += blockDim.x) {
+      // load into smem
+      {
+        const int j = offset + tid;
+        if(j < stride) {
+          smem[tid] = softmax[strideOffset + j];
+          smem[tid + blockDim.x] = upstreamGrad[strideOffset + j];
+        }
+        __syncthreads();
+      }
+
+
+      for(int k = 0; k < blockDim.x; k++) {
+        const int j = offset + k;
+        if(j < stride) {
+          ftype yj = smem[k];
+          ftype gj = smem[k + blockDim.x];
+
+          auto jacobian = (i == j) ? yi * (1 - yj) : -yi * yj;
+          grad += gj * jacobian;
+        }
+      }
+      __syncthreads();
+    }
+
+    if(isNotPadded) {
+      res[gid] = grad;
+    } 
+  }
+}
+
+namespace cuda_impl {
+  void reluBackward(Tensor& res, const Tensor& upstreamGrad, const Tensor& parent) {
+    constexpr int threadsPerBlock = 256;
+    const int blocks = (upstreamGrad.getSize() + threadsPerBlock - 1) / threadsPerBlock;
+
+    reluBackwardKernel<<<blocks, threadsPerBlock>>>(res.getData(), upstreamGrad.getData(), parent.getData(), res.getSize());
+    cudaErrchk(cudaDeviceSynchronize());
+  }
+
+  void leakyReluBackward(Tensor& res, const Tensor& upstreamGrad, const Tensor& parent, ftype eps) {
+    constexpr int threadsPerBlock = 256;
+    const int blocks = (upstreamGrad.getSize() + threadsPerBlock - 1) / threadsPerBlock;
+
+    leakyReluBackwardKernel<<<blocks, threadsPerBlock>>>(res.getData(), upstreamGrad.getData(), parent.getData(), eps, res.getSize());
+    cudaErrchk(cudaDeviceSynchronize());
+  }
+
+  void sigmoidBackward(Tensor& res, const Tensor& upstreamGrad, const Tensor& sigmoid) {
+    constexpr int threadsPerBlock = 256;
+    const int blocks = (upstreamGrad.getSize() + threadsPerBlock - 1) / threadsPerBlock;
+
+    sigmoidBackwardKernel<<<blocks, threadsPerBlock>>>(res.getData(), upstreamGrad.getData(), sigmoid.getData(), res.getSize());
+    cudaErrchk(cudaDeviceSynchronize());
+  }
+
+  /**
+   * @brief The backward of the softmax. Due to optimization this function distinguishes three cases of stride size, where stride
+   * is the size of the dimension the softmax operation is applied to. The two cases are a stride either fitting into one block or not.
+   */
+  void softmaxBackward(Tensor& res, const Tensor& upstreamGrad, const Tensor& softmax) {
+    assert(upstreamGrad.getSize() == softmax.getSize());
+
+    constexpr int maxThreadsPerBlock = 256;
+    const int stride = softmax.getDims()[-1];
+    
+    if(stride < maxThreadsPerBlock) {
+      const int threadsPerStride = max(1, ((stride + 31) / 32)) * 32; // == warps per stride * 32
+
+      // min over maximum possible strides per block and actual number of strides
+      const int stridesPerBlock = min(maxThreadsPerBlock / threadsPerStride, softmax.getSize() / stride);
+      const int strideWidthPerBlock = stridesPerBlock * stride; // for smem idx computation
+      
+      int threadsPerBlock = 1;
+      while(threadsPerBlock < threadsPerStride * stridesPerBlock) threadsPerBlock <<= 1; 
+      // threadsPerBlock now larger than threadsPerStride * stridesPerBlock
+      const int blocks = (upstreamGrad.getSize() + threadsPerBlock - 1) / threadsPerBlock;
+
+      softmaxBackwardKernelOneBlock<<<blocks, threadsPerBlock, 2 * strideWidthPerBlock * sizeof(ftype)>>>(
+          res.getData(), upstreamGrad.getData(), softmax.getData(), stride, strideWidthPerBlock, threadsPerStride, softmax.getSize());
+    }
+    else {
+      constexpr int maxThreadsPerBlock = 256; 
+
+      const int nStrides = softmax.getSize() / stride;
+      const int threadsPerBlock = maxThreadsPerBlock; // TODO: do that one better, this can result in gross imbalance; also for normal softmax
+      const int blocksPerStride = (stride + threadsPerBlock - 1) / threadsPerBlock; 
+
+      softmaxBackwardKernelLargePass<<<blocksPerStride * nStrides, threadsPerBlock, 2 * threadsPerBlock * sizeof(ftype)>>>(
+                                       res.getData(), upstreamGrad.getData(), softmax.getData(), blocksPerStride, stride);
+    }
+    cudaErrchk(cudaDeviceSynchronize());
+  }
+}
diff --git a/src/backend/computational_graph/activation_functions/cuda/activation_nodes.cuh b/src/backend/computational_graph/activation_functions/cuda/activation_nodes.cuh
new file mode 100644
index 0000000..44624e9
--- /dev/null
+++ b/src/backend/computational_graph/activation_functions/cuda/activation_nodes.cuh
@@ -0,0 +1,27 @@
+/**
+ * @file activation_nodes.cuh
+ * @author Robert Baumgartner (r.baumgartner-1@tudelft.nl)
+ * @brief
+ * @version 0.1
+ * @date 2026-03-23
+ *
+ * @copyright Copyright (c) 2026
+ *
+ */
+
+#pragma once
+
+#ifndef __CUDA
+static_assert(false, "File should not be included without CUDA enabled");
+#endif // __CUDA
+
+#include "utility/global_params.h"
+#include "data_modeling/tensor.h"
+
+namespace cuda_impl {
+  void reluBackward(Tensor& res, const Tensor& upstreamGrad, const Tensor& parent);
+  void leakyReluBackward(Tensor& res, const Tensor& upstreamGrad, const Tensor& parent, ftype eps);
+
+  void sigmoidBackward(Tensor& res, const Tensor& upstreamGrad, const Tensor& sigmoid);
+  void softmaxBackward(Tensor& res, const Tensor& upstreamGrad, const Tensor& softmax);
+}
diff --git a/src/backend/computational_graph/activation_functions/leaky_relu_node.cpp b/src/backend/computational_graph/activation_functions/leaky_relu_node.cpp
index 83de4ca..24607e8 100644
--- a/src/backend/computational_graph/activation_functions/leaky_relu_node.cpp
+++ b/src/backend/computational_graph/activation_functions/leaky_relu_node.cpp
@@ -1,31 +1,49 @@
 /**
  * @file leaky_relu_node.cpp
  * @author Robert Baumgartner (r.baumgartner-1@tudelft.nl)
- * @brief 
+ * @brief
  * @version 0.1
  * @date 2026-03-07
- * 
+ *
  * @copyright Copyright (c) 2026
- * 
+ *
  */
 
 #include "leaky_relu_node.h"
 
 #include <utility>
 
+#ifdef __CUDA
+#include "computational_graph/activation_functions/cuda/activation_nodes.cuh"
+#else
+#include <stdexcept>
+#endif
+
 using namespace std;
 using namespace cgraph;
 
 vector<shared_ptr<Tensor>> LeakyReLuNode::backward(const Tensor& upstreamGrad) {
   assert(!upstreamGrad.getRequiresGrad());
-  constexpr ftype zero = 0.0;
-  
+
   auto res = make_shared<Tensor>(upstreamGrad.getDims(), upstreamGrad.getDevice(), false);
-  
   const auto& parent = parents[0];
-  for(tensorSize_t i=0; i<upstreamGrad.getSize(); i++){
-    res->set((*parent)[i] > zero ? upstreamGrad[i] : upstreamGrad[i] * eps, i);
+
+  switch(upstreamGrad.getDevice()) {
+    case Device::CPU: {
+      constexpr ftype zero = 0.0;
+      for(tensorSize_t i = 0; i < upstreamGrad.getSize(); i++){
+        res->set((*parent)[i] > zero ? upstreamGrad[i] : upstreamGrad[i] * eps, i);
+      }
+      break;
+    }
+    case Device::CUDA:
+    #ifdef __CUDA
+      cuda_impl::leakyReluBackward(*res, upstreamGrad, *parent, eps);
+    #else
+      __throw_invalid_argument("Attempted to give CUDA tensor");
+    #endif
+      break;
   }
 
   return {res};
-}
\ No newline at end of file
+}
diff --git a/src/backend/computational_graph/activation_functions/relu_node.cpp b/src/backend/computational_graph/activation_functions/relu_node.cpp
index 3fcc958..78c6ec3 100644
--- a/src/backend/computational_graph/activation_functions/relu_node.cpp
+++ b/src/backend/computational_graph/activation_functions/relu_node.cpp
@@ -1,31 +1,49 @@
 /**
  * @file relu_node.cpp
  * @author Robert Baumgartner (r.baumgartner-1@tudelft.nl)
- * @brief 
+ * @brief
  * @version 0.1
  * @date 2026-02-15
- * 
+ *
  * @copyright Copyright (c) 2026
- * 
+ *
  */
 
 #include "relu_node.h"
 
 #include <utility>
 
+#ifdef __CUDA
+#include "computational_graph/activation_functions/cuda/activation_nodes.cuh"
+#else
+#include <stdexcept>
+#endif
+
 using namespace std;
 using namespace cgraph;
 
 vector<shared_ptr<Tensor>> ReLuNode::backward(const Tensor& upstreamGrad) {
   assert(!upstreamGrad.getRequiresGrad());
-  constexpr ftype zero = 0.0;
-  
-  auto res = make_shared<Tensor>(upstreamGrad.getDims(), upstreamGrad.getDevice(), false);
 
+  auto res = make_shared<Tensor>(upstreamGrad.getDims(), upstreamGrad.getDevice(), false);
   const auto& parent = parents[0];
-  for(tensorSize_t i=0; i<upstreamGrad.getSize(); i++){
-    res->set((*parent)[i] > zero ? upstreamGrad[i] : zero, i);
+
+  switch(upstreamGrad.getDevice()) {
+    case Device::CPU: {
+      constexpr ftype zero = 0.0;
+      for(tensorSize_t i = 0; i < upstreamGrad.getSize(); i++){
+        res->set((*parent)[i] > zero ? upstreamGrad[i] : zero, i);
+      }
+      break;
+    }
+    case Device::CUDA:
+    #ifdef __CUDA
+      cuda_impl::reluBackward(*res, upstreamGrad, *parent);
+    #else
+      __throw_invalid_argument("Attempted to give CUDA tensor");
+    #endif
+      break;
   }
 
   return {res};
-}
\ No newline at end of file
+}
diff --git a/src/backend/computational_graph/activation_functions/sigmoid_node.cpp b/src/backend/computational_graph/activation_functions/sigmoid_node.cpp
index 5873724..7f6f85b 100644
--- a/src/backend/computational_graph/activation_functions/sigmoid_node.cpp
+++ b/src/backend/computational_graph/activation_functions/sigmoid_node.cpp
@@ -1,35 +1,51 @@
 /**
  * @file sigmoid_node.cpp
  * @author Robert Baumgartner (r.baumgartner-1@tudelft.nl)
- * @brief 
+ * @brief
  * @version 0.1
  * @date 2026-03-14
- * 
+ *
  * @copyright Copyright (c) 2026
- * 
+ *
  */
 
 #include "sigmoid_node.h"
 
 #include <utility>
 
+#ifdef __CUDA
+#include "computational_graph/activation_functions/cuda/activation_nodes.cuh"
+#else
+#include <stdexcept>
+#endif
+
 using namespace std;
 using namespace cgraph;
 
 vector<shared_ptr<Tensor>> SigmoidNode::backward(const Tensor& upstreamGrad) {
   assert(!upstreamGrad.getRequiresGrad());
-  constexpr ftype zero = 0.0;
-  
+
   auto res = make_shared<Tensor>(upstreamGrad.getDims(), upstreamGrad.getDevice(), false);
 
-  // s is result from forward pass sigmoid
-  auto derivative = [](ftype s){
-    return s * (1-s);
-  };
+  switch(upstreamGrad.getDevice()) {
+    case Device::CPU: {
+      auto derivative = [](ftype s){
+        return s * (1 - s);
+      };
 
-  for(tensorSize_t i=0; i<upstreamGrad.getSize(); i++){
-    res->set(derivative((*sigmoid)[i]) * upstreamGrad[i], i);
+      for(tensorSize_t i = 0; i < upstreamGrad.getSize(); i++){
+        res->set(derivative((*sigmoid)[i]) * upstreamGrad[i], i);
+      }
+      break;
+    }
+    case Device::CUDA:
+    #ifdef __CUDA
+      cuda_impl::sigmoidBackward(*res, upstreamGrad, *sigmoid);
+    #else
+      __throw_invalid_argument("Attempted to give CUDA tensor");
+    #endif
+      break;
   }
 
   return {res};
-}
\ No newline at end of file
+}
diff --git a/src/backend/computational_graph/activation_functions/softmax_node.cpp b/src/backend/computational_graph/activation_functions/softmax_node.cpp
index 8603355..57e2dbe 100644
--- a/src/backend/computational_graph/activation_functions/softmax_node.cpp
+++ b/src/backend/computational_graph/activation_functions/softmax_node.cpp
@@ -1,12 +1,12 @@
 /**
  * @file softmax_node.cpp
  * @author Robert Baumgartner (r.baumgartner-1@tudelft.nl)
- * @brief 
+ * @brief
  * @version 0.1
  * @date 2026-03-15
- * 
+ *
  * @copyright Copyright (c) 2026
- * 
+ *
  */
 
 #include "softmax_node.h"
@@ -15,31 +15,52 @@
 
 #include <iostream>
 
+#ifdef __CUDA
+#include "computational_graph/activation_functions/cuda/activation_nodes.cuh"
+#else
+#include <stdexcept>
+#endif
+
 using namespace std;
 using namespace cgraph;
 
+/**
+ * @brief Shape assumptions of softmax forward: [bsize, dim1, dim2, ..., stride]
+ */
 vector< shared_ptr<Tensor> > SoftmaxNode::backward(const Tensor& upstreamGrad) {
   assert(!upstreamGrad.getRequiresGrad());
-    
+
   const auto& yPred = parents[0];
   auto res = make_shared<Tensor>(yPred->createEmptyCopy());
 
-  const auto bSize = yPred->getDims()[0];
-  assert(bSize>0);
-
-  for(tensorDim_t b=0; b<bSize; b++){
-    for(tensorDim_t i=0; i<yPred->getDims()[1]; i++){
-      ftype grad = 0;
-      const ftype yi = softmax->get(b, i);
-      
-      for(tensorDim_t j=0; j<yPred->getDims()[1]; j++){
-        ftype yj = softmax->get(b, j);
-        ftype jacobian = (i==j) ? yi*(1-yj) : -yi*yj;
-        grad += upstreamGrad.get(b, j) * jacobian;
+  switch(upstreamGrad.getDevice()) {
+    case Device::CPU: {
+      const tensorSize_t stride = yPred->getDims()[-1];
+      tensorSize_t offset = 0;
+      while(offset < yPred->getSize()) {
+        for(tensorSize_t i = 0; i < stride; i++) {
+          ftype grad = 0;
+          const ftype yi = softmax->get(offset + i);
+
+          for(tensorSize_t j = 0; j < stride; j++) {
+            const ftype yj = softmax->get(offset + j);
+            const ftype jacobian = (i == j) ? yi * (1 - yj) : -yi * yj;
+            grad += upstreamGrad.get(offset + j) * jacobian;
+          }
+          res->set(grad, offset + i);
+        }
+        offset += stride;
       }
-      res->set(grad, b, i);
+      break;
     }
+    case Device::CUDA:
+    #ifdef __CUDA
+      cuda_impl::softmaxBackward(*res, upstreamGrad, *softmax);
+    #else
+      __throw_invalid_argument("Attempted to give CUDA tensor");
+    #endif
+      break;
   }
 
   return {res};
-}
\ No newline at end of file
+}
diff --git a/src/backend/computational_graph/loss_functions/bce_node.cpp b/src/backend/computational_graph/loss_functions/bce_node.cpp
index add016f..cc09a44 100644
--- a/src/backend/computational_graph/loss_functions/bce_node.cpp
+++ b/src/backend/computational_graph/loss_functions/bce_node.cpp
@@ -1,18 +1,23 @@
 /**
  * @file bce_node.cpp
  * @author Robert Baumgartner (r.baumgartner-1@tudelft.nl)
- * @brief 
+ * @brief
  * @version 0.1
  * @date 2026-03-14
- * 
+ *
  * @copyright Copyright (c) 2026
- * 
+ *
  */
 
 #include "bce_node.h"
-
 #include "data_modeling/tensor_functions.h"
 
+#ifdef __CUDA
+#include "computational_graph/loss_functions/cuda/loss_nodes.cuh"
+#else
+#include <stdexcept>
+#endif
+
 using namespace std;
 using namespace cgraph;
 
@@ -22,14 +27,27 @@ vector< shared_ptr<Tensor> > BceNode::backward(const Tensor& upstreamGrad) {
   const auto& yPred = parents[0];
   auto res = make_shared<Tensor>(yPred->createEmptyCopy());
 
-  ftype bSize = yPred->getDims()[0];
-  for(tensorSize_t i=0; i<yPred->getDims()[0]; i++){
-    auto yi = (*yTrue)[i];
-    auto yiHat = (*yPred)[i];
+  switch(upstreamGrad.getDevice()) {
+    case Device::CPU: {
+      const ftype bSize = yPred->getDims()[0];
 
-    auto g = -yi/std::max(yiHat, epsBce) + (1-yi)/std::max(1-yiHat, epsBce);
-    res->set(g/bSize, i);
+      for(tensorSize_t i = 0; i < yPred->getDims()[0]; i++){
+        const auto yi = (*yTrue)[i];
+        const auto yiHat = (*yPred)[i];
+        
+        const auto g = -yi / std::max(yiHat, EPS_BCE) + (1 - yi) / std::max(1 - yiHat, EPS_BCE);
+        res->set(g / bSize, i);
+      }
+      break;
+    }
+    case Device::CUDA:
+    #ifdef __CUDA
+      cuda_impl::bceBackward(*res, *yPred, *yTrue);
+    #else
+      __throw_invalid_argument("Attempted to give CUDA tensor");
+    #endif
+      break;
   }
-  
+
   return {res};
-}
\ No newline at end of file
+}
diff --git a/src/backend/computational_graph/loss_functions/bce_sigmoid_node.cpp b/src/backend/computational_graph/loss_functions/bce_sigmoid_node.cpp
index 998e110..2c1420b 100644
--- a/src/backend/computational_graph/loss_functions/bce_sigmoid_node.cpp
+++ b/src/backend/computational_graph/loss_functions/bce_sigmoid_node.cpp
@@ -1,46 +1,62 @@
 /**
  * @file bce_sigmoid_node.cpp
  * @author Robert Baumgartner (r.baumgartner-1@tudelft.nl)
- * @brief 
+ * @brief
  * @version 0.1
  * @date 2026-03-17
- * 
+ *
  * @copyright Copyright (c) 2026
- * 
+ *
  */
 
 #include "bce_sigmoid_node.h"
-
 #include "data_modeling/tensor_functions.h"
 
 #include <cmath>
 
+#ifdef __CUDA
+#include "computational_graph/loss_functions/cuda/loss_nodes.cuh"
+#else
+#include <stdexcept>
+#endif
+
 using namespace std;
 using namespace cgraph;
 
 vector< shared_ptr<Tensor> > BceSigmoidNode::backward(const Tensor& upstreamGrad) {
   assert(!upstreamGrad.getRequiresGrad());
-  
-  auto sigmoid = [](ftype x){
-    constexpr ftype one = 1.0;
-    if(x>=0){
-      return one / (one + exp(-x));
-    }
-    auto e = exp(x);
-    return e / (one + e);
-  };
 
   const auto& logits = parents[0];
   auto res = make_shared<Tensor>(logits->createEmptyCopy());
 
-  ftype bSize = logits->getDims()[0];
-  for(tensorSize_t i=0; i<logits->getDims()[0]; i++){
-    auto y = (*yTrue)[i];
-    auto s = sigmoid((*logits)[i]);
-
-    auto g = s - y;
-    res->set(g/bSize, i);
+  switch(upstreamGrad.getDevice()) {
+    case Device::CPU: {
+      auto sigmoid = [](ftype x){
+        constexpr ftype one = 1.0;
+        if(x >= 0) {
+          return one / (one + exp(-x));
+        }
+        auto e = exp(x);
+        return e / (one + e);
+      };
+
+      const ftype bSize = logits->getDims()[0];
+      for(tensorSize_t i = 0; i < logits->getDims()[0]; i++){
+        auto y = (*yTrue)[i];
+        auto s = sigmoid((*logits)[i]);
+        auto g = s - y;
+        res->set(g / bSize, i);
+      }
+      break;
+    }
+    case Device::CUDA:
+    #ifdef __CUDA
+      cuda_impl::bceSigmoidBackward(*res, *logits, *yTrue);
+    #else
+      __throw_invalid_argument("Attempted to give CUDA tensor");
+    #endif
+      break;
   }
-  
+
   return {res};
-}
\ No newline at end of file
+}
diff --git a/src/backend/computational_graph/loss_functions/crossentropy_node.cpp b/src/backend/computational_graph/loss_functions/crossentropy_node.cpp
index 249de43..be15b0f 100644
--- a/src/backend/computational_graph/loss_functions/crossentropy_node.cpp
+++ b/src/backend/computational_graph/loss_functions/crossentropy_node.cpp
@@ -1,37 +1,54 @@
 /**
- * @file add_node.cpp
+ * @file crossentropy_node.cpp
  * @author Robert Baumgartner (r.baumgartner-1@tudelft.nl)
- * @brief 
+ * @brief
  * @version 0.1
  * @date 2026-02-03
- * 
+ *
  * @copyright Copyright (c) 2026
- * 
+ *
  */
 
 #include "crossentropy_node.h"
-
 #include "data_modeling/tensor_functions.h"
 
+#ifdef __CUDA
+#include "computational_graph/loss_functions/cuda/loss_nodes.cuh"
+#else
+#include <stdexcept>
+#endif
+
 using namespace std;
 using namespace cgraph;
 
+/**
+ * @brief Backward function on crossentropy-node. Uses cached values of forward pass 
+ * for higher efficiency.
+ */
 vector< shared_ptr<Tensor> > CrossEntropyNode::backward(const Tensor& upstreamGrad) {
   assert(!upstreamGrad.getRequiresGrad());
-  
+
   const auto& yPred = parents[0];
   auto res = make_shared<Tensor>(yPred->createEmptyCopy());
 
-  ftype bSize = yPred->getDims()[0];
-  for(tensorDim_t i=0; i<yPred->getDims()[0]; i++){
-    for(tensorDim_t j=0; j<yPred->getDims()[1]; j++){
-      auto yij = yTrue->get(i, j);
-      auto yijHat = yPred->get(i, j);
-
-      auto g = -yij/std::max(yijHat, epsCrossentropy);
-      res->set(g/bSize, i, j);
+  switch(upstreamGrad.getDevice()) {
+    case Device::CPU: {
+      const tensorSize_t stride = yPred->getDims()[-1];
+      const ftype bSize = static_cast<ftype>(yPred->getSize() / stride);
+      for(tensorSize_t i = 0; i < yPred->getSize(); i++) {
+        auto g = -(*yTrue)[i] / std::max((*yPred)[i], EPS_CROSSENTROPY);
+        res->set(g / bSize, i);
+      }
+      break;
     }
+    case Device::CUDA:
+    #ifdef __CUDA
+      cuda_impl::crossEntropyBackward(*res, *yPred, *yTrue);
+    #else
+      __throw_invalid_argument("Attempted to give CUDA tensor");
+    #endif
+      break;
   }
-  
+
   return {res};
-}
\ No newline at end of file
+}
diff --git a/src/backend/computational_graph/loss_functions/crossentropy_softmax_node.cpp b/src/backend/computational_graph/loss_functions/crossentropy_softmax_node.cpp
index ca0d1c7..0ab4356 100644
--- a/src/backend/computational_graph/loss_functions/crossentropy_softmax_node.cpp
+++ b/src/backend/computational_graph/loss_functions/crossentropy_softmax_node.cpp
@@ -1,37 +1,58 @@
 /**
  * @file crossentropy_softmax_node.cpp
  * @author Robert Baumgartner (r.baumgartner-1@tudelft.nl)
- * @brief 
+ * @brief
  * @version 0.1
  * @date 2026-03-17
- * 
+ *
  * @copyright Copyright (c) 2026
- * 
+ *
  */
 
 #include "crossentropy_softmax_node.h"
-
 #include "module/activation_functions/softmax.h"
 
+#ifdef __CUDA
+#include "computational_graph/loss_functions/cuda/loss_nodes.cuh"
+#else
+#include <stdexcept>
+#endif
+
 using namespace std;
 using namespace cgraph;
 
 vector< shared_ptr<Tensor> > CrossEntropySoftmaxNode::backward(const Tensor& upstreamGrad) {
-  assert(!upstreamGrad.getRequiresGrad());
+  assert(!upstreamGrad.getRequiresGrad() && logits.nDims() == 2);
 
   const auto& logits = parents[0];
   auto res = make_shared<Tensor>(logits->createEmptyCopy());
 
-  const auto softmax = module::Softmax();
-  const auto s = softmax(*logits);
-
-  ftype bSize = logits->getDims()[0];
-  for(tensorSize_t b=0; b<logits->getDims()[0]; b++){
-    for(tensorSize_t i=0; i<logits->getDims()[1]; i++){
-      auto g = s.get(b, i) - yTrue->get(b, i);
-      res->set(g / bSize, b, i);
+  switch(upstreamGrad.getDevice()) {
+    case Device::CPU: 
+    {
+      static const auto softmax = module::Softmax();
+      const auto s = softmax(*logits);
+
+      const tensorSize_t stride = logits->getDims().get(-1);
+      const tensorSize_t bSize = logits->getSize() / stride;
+
+      for(tensorSize_t b = 0; b < bSize; b++){
+        for(tensorSize_t i = 0; i < stride; i++){
+          const tensorSize_t flatIdx = b * stride + i;
+          auto g = s[flatIdx] - (*yTrue)[flatIdx];
+          res->set(g / bSize, flatIdx);
+        }
+      }
+      break;
     }
+    case Device::CUDA:
+    #ifdef __CUDA
+      cuda_impl::crossEntropySoftmaxBackward(*res, *logits, *yTrue);
+    #else
+      __throw_invalid_argument("Attempted to give CUDA tensor");
+    #endif
+      break;
   }
 
   return {res};
-}
\ No newline at end of file
+}
diff --git a/src/backend/computational_graph/loss_functions/cuda/loss_nodes.cu b/src/backend/computational_graph/loss_functions/cuda/loss_nodes.cu
new file mode 100644
index 0000000..e2b0508
--- /dev/null
+++ b/src/backend/computational_graph/loss_functions/cuda/loss_nodes.cu
@@ -0,0 +1,162 @@
+/**
+ * @file loss_nodes.cu
+ * @author Robert Baumgartner (r.baumgartner-1@tudelft.nl)
+ * @brief
+ * @version 0.1
+ * @date 2026-03-23
+ *
+ * @copyright Copyright (c) 2026
+ *
+ */
+
+#ifndef __CUDA
+static_assert(false, "File should not be compiled without CUDA enabled");
+#endif // __CUDA
+
+#include "loss_nodes.cuh"
+#include "module/activation_functions/softmax.h"
+
+#include "shared/cuda/common_kernels.cuh"
+#include "utility/cuda/cuda_common.cuh"
+#include "utility/global_params.h"
+
+#include <type_traits>
+
+using namespace std;
+
+namespace {
+  using namespace cuda_impl;
+  
+  /**
+   * @brief Does what you think it does.
+   */
+  __global__ void bceBackwardKernel(ftype* const res, const ftype* const yPred, const ftype* const yTrue, const ftype bSize, const tensorSize_t size) {
+    const int gid = blockIdx.x * blockDim.x + threadIdx.x;
+    if(gid >= size) {
+      return;
+    }
+
+    const auto yi = yTrue[gid];
+    const auto yiHat = yPred[gid];
+
+    const auto g = -yi / cudaMax<ftype>(yiHat, EPS_BCE) + (1 - yi) / cudaMax<ftype>(1 - yiHat, EPS_BCE);
+    res[gid] = g / bSize;
+  }
+
+  /**
+   * @brief BCE backward kernel appplied with sigmoid.
+   * 
+   * @param bSize Batch-size.
+   * @param sigmoids Sigmoids from forward pass.
+   */
+  __global__ void bceSigmoidBackwardKernel(ftype* const res, const ftype* const logits, const ftype* const yTrue, const ftype bSize, const tensorSize_t size) {
+    const int gid = blockIdx.x * blockDim.x + threadIdx.x;
+    if(gid >= size) {
+      return;
+    }
+
+    const auto y = yTrue[gid];
+    const auto s = cudaSigmoid<ftype>(logits[gid]);
+
+    const auto g = s - y;
+    res[gid] = g / bSize;
+  }
+
+  /**
+   * @brief The simple cross-entropy backward kernel.
+   * 
+   * @param bSize The batch-size.
+   */
+  __global__ void crossEntropyBackwardKernel(ftype* const res, const ftype* const yPred, const ftype* const yTrue, const ftype nSamples, const tensorSize_t size) {
+    const int gid = blockIdx.x * blockDim.x + threadIdx.x;
+    if(gid >= size) {
+      return;
+    }
+
+    const auto g = -yTrue[gid] / cudaMax<ftype>(yPred[gid], EPS_CROSSENTROPY);
+    res[gid] = g / nSamples;
+  }
+
+  /**
+   * @brief Does what you think it does.
+   */
+  __global__ void crossEntropySoftmaxBackwardKernel(ftype* const res, const ftype* const softmaxedLogits, const ftype* const yTrue, 
+                                                    const ftype nSamples, const tensorSize_t size) {
+    const int gid = blockIdx.x * blockDim.x + threadIdx.x;
+    if(gid >= size) {
+      return;
+    }
+
+    res[gid] = (softmaxedLogits[gid] - yTrue[gid]) / nSamples;
+  }
+
+  /**
+   * @brief RMSE backward kernel. bSize = batch-size, rmse = the rmse from the forward pass.
+   */
+  __global__ void rmseBackwardKernel(ftype* const res, const ftype* const yPred, const ftype* const yTrue,
+                                     const ftype rmse, const ftype bSize, const tensorSize_t size) {
+    const int gid = blockIdx.x * blockDim.x + threadIdx.x;
+    if(gid >= size) {
+      return;
+    }
+    
+    const ftype yi = yTrue[gid];
+    const ftype yiHat = yPred[gid];
+
+    const ftype denom = rmse * bSize + EPS_RMSE;
+    const ftype g = (yiHat - yi) / denom;
+
+    res[gid] = g;
+  }
+}
+
+namespace cuda_impl {
+  void bceBackward(Tensor& res, const Tensor& yPred, const Tensor& yTrue) {
+    constexpr int threadsPerBlock = 256;
+    const int blocks = (yPred.getSize() + threadsPerBlock - 1) / threadsPerBlock;
+
+    bceBackwardKernel<<<blocks, threadsPerBlock>>>(res.getData(), yPred.getData(), yTrue.getData(), yPred.getDims()[0], yTrue.getSize());
+    cudaErrchk(cudaDeviceSynchronize());
+  }
+
+  void bceSigmoidBackward(Tensor& res, const Tensor& logits, const Tensor& yTrue) {
+    constexpr int threadsPerBlock = 256;
+    const int blocks = (logits.getSize() + threadsPerBlock - 1) / threadsPerBlock;
+
+    bceSigmoidBackwardKernel<<<blocks, threadsPerBlock>>>(res.getData(), logits.getData(), yTrue.getData(), logits.getDims()[0], logits.getSize());
+    cudaErrchk(cudaDeviceSynchronize());
+  }
+
+  void crossEntropyBackward(Tensor& res, const Tensor& yPred, const Tensor& yTrue) {
+    constexpr int threadsPerBlock = 256;
+    const int blocks = (yPred.getSize() + threadsPerBlock - 1) / threadsPerBlock;
+
+    const tensorSize_t stride = yPred.getDims()[-1];
+    const ftype nSamples = static_cast<ftype>(yPred.getSize() / stride);
+
+    crossEntropyBackwardKernel<<<blocks, threadsPerBlock>>>(res.getData(), yPred.getData(), yTrue.getData(), nSamples, yTrue.getSize());
+    cudaErrchk(cudaDeviceSynchronize());
+  }
+
+  void crossEntropySoftmaxBackward(Tensor& res, const Tensor& logits, const Tensor& yTrue) {
+    constexpr int threadsPerBlock = 256;
+    const int blocks = (logits.getSize() + threadsPerBlock - 1) / threadsPerBlock;
+
+    static const auto softmax = module::Softmax();
+    const auto softmaxedLogits = softmax(logits);
+
+    const tensorSize_t stride = logits.getDims().get(-1);
+    const auto nSamples = static_cast<ftype>(logits.getSize() / stride);
+
+    crossEntropySoftmaxBackwardKernel<<<blocks, threadsPerBlock>>>(res.getData(), softmaxedLogits.getData(), yTrue.getData(), nSamples, logits.getSize());
+    cudaErrchk(cudaDeviceSynchronize());
+  }
+
+  void rmseBackward(Tensor& res, const Tensor& yPred, const Tensor& yTrue, ftype rmse) {
+    constexpr int threadsPerBlock = 256;
+    const int blocks = (yPred.getSize() + threadsPerBlock - 1) / threadsPerBlock;
+
+    rmseBackwardKernel<<<blocks, threadsPerBlock>>>(res.getData(), yPred.getData(), yTrue.getData(), rmse, static_cast<ftype>(yPred.getDims()[0]), yPred.getSize());
+    cudaErrchk(cudaDeviceSynchronize());
+  }
+}
diff --git a/src/backend/computational_graph/loss_functions/cuda/loss_nodes.cuh b/src/backend/computational_graph/loss_functions/cuda/loss_nodes.cuh
new file mode 100644
index 0000000..24486bf
--- /dev/null
+++ b/src/backend/computational_graph/loss_functions/cuda/loss_nodes.cuh
@@ -0,0 +1,29 @@
+/**
+ * @file loss_nodes.cuh
+ * @author Robert Baumgartner (r.baumgartner-1@tudelft.nl)
+ * @brief
+ * @version 0.1
+ * @date 2026-03-23
+ *
+ * @copyright Copyright (c) 2026
+ *
+ */
+
+#pragma once
+
+#ifndef __CUDA
+static_assert(false, "File should not be included without CUDA enabled");
+#endif // __CUDA
+
+#include "utility/global_params.h"
+#include "data_modeling/tensor.h"
+
+namespace cuda_impl {
+  void bceBackward(Tensor& res, const Tensor& yPred, const Tensor& yTrue);
+  void bceSigmoidBackward(Tensor& res, const Tensor& logits, const Tensor& yTrue);
+
+  void crossEntropyBackward(Tensor& res, const Tensor& yPred, const Tensor& yTrue);
+  void crossEntropySoftmaxBackward(Tensor& res, const Tensor& logits, const Tensor& yTrue);
+
+  void rmseBackward(Tensor& res, const Tensor& yPred, const Tensor& yTrue, ftype rmse);
+}
diff --git a/src/backend/computational_graph/loss_functions/rmse_node.cpp b/src/backend/computational_graph/loss_functions/rmse_node.cpp
index e3eb11e..b171ef7 100644
--- a/src/backend/computational_graph/loss_functions/rmse_node.cpp
+++ b/src/backend/computational_graph/loss_functions/rmse_node.cpp
@@ -1,39 +1,57 @@
 /**
  * @file rmse_node.cpp
  * @author Robert Baumgartner (r.baumgartner-1@tudelft.nl)
- * @brief 
+ * @brief
  * @version 0.1
  * @date 2026-03-14
- * 
+ *
  * @copyright Copyright (c) 2026
- * 
+ *
  */
 
 #include "rmse_node.h"
-
 #include "data_modeling/tensor_functions.h"
 
 #include <iostream>
 
+#ifdef __CUDA
+#include "computational_graph/loss_functions/cuda/loss_nodes.cuh"
+#else
+#include <stdexcept>
+#endif
+
 using namespace std;
 using namespace cgraph;
 
 vector< shared_ptr<Tensor> > RmseNode::backward(const Tensor& upstreamGrad) {
   assert(!upstreamGrad.getRequiresGrad());
-  constexpr ftype eps = 1e-9;
 
   const auto& yPred = parents[0];
   auto res = make_shared<Tensor>(yPred->createEmptyCopy());
 
-  ftype bSize = yPred->getDims()[0];
-  for(tensorSize_t i=0; i<yPred->getDims()[0]; i++){
-    auto yi = (*yTrue)[i];
-    auto yiHat = (*yPred)[i];
-
-    auto denom = rmse * bSize + eps;
-    auto g = (yiHat-yi) / denom;
-    res->set(g, i);
+  switch(upstreamGrad.getDevice()) {
+    case Device::CPU: 
+    {
+      const ftype bSize = yPred->getDims()[0];
+      for(tensorSize_t i = 0; i < yPred->getSize(); i++){
+        auto yi = (*yTrue)[i];
+        auto yiHat = (*yPred)[i];
+
+        auto denom = rmse * bSize + EPS_RMSE;
+        auto g = (yiHat - yi) / denom;
+        
+        res->set(g, i);
+      }
+      break;
+    }
+    case Device::CUDA:
+    #ifdef __CUDA
+      cuda_impl::rmseBackward(*res, *yPred, *yTrue, rmse);
+    #else
+      __throw_invalid_argument("Attempted to give CUDA tensor");
+    #endif
+      break;
   }
-  
+
   return {res};
-}
\ No newline at end of file
+}
diff --git a/src/backend/computational_graph/tensor_ops/cuda/tensor_ops_nodes.cu b/src/backend/computational_graph/tensor_ops/cuda/tensor_ops_nodes.cu
new file mode 100644
index 0000000..d473ed4
--- /dev/null
+++ b/src/backend/computational_graph/tensor_ops/cuda/tensor_ops_nodes.cu
@@ -0,0 +1,45 @@
+/**
+ * @file tensor_ops_nodes.cu
+ * @author Robert Baumgartner (r.baumgartner-1@tudelft.nl)
+ * @brief
+ * @version 0.1
+ * @date 2026-05-20
+ *
+ * @copyright Copyright (c) 2026
+ *
+ */
+
+#ifndef __CUDA
+static_assert(false, "File should not be compiled without CUDA enabled");
+#endif // __CUDA
+
+#include "tensor_ops_nodes.cuh"
+#include "utility/cuda/cuda_common.cuh"
+
+using namespace std;
+
+namespace {
+  __global__ void scalarMulKernel(ftype* const res, const ftype* const upstreamGrad, const ftype factor, const tensorSize_t size) {
+    const int gid = blockIdx.x * blockDim.x + threadIdx.x;
+    if(gid >= size) return;
+    res[gid] = upstreamGrad[gid] * factor;
+  }
+}
+
+namespace cuda_impl {
+  void scalarMulBackward(Tensor& res, const Tensor& upstreamGrad, ftype factor) {
+    constexpr int threadsPerBlock = 256;
+    const int blocks = (upstreamGrad.getSize() + threadsPerBlock - 1) / threadsPerBlock;
+
+    scalarMulKernel<<<blocks, threadsPerBlock>>>(res.getData(), upstreamGrad.getData(), factor, upstreamGrad.getSize());
+    cudaErrchk(cudaDeviceSynchronize());
+  }
+
+  void getterBackward(Tensor& res, ftype val, tensorSize_t linearIdx) {
+    cudaErrchk(cudaMemset(res.getData(), 0, res.getSize() * sizeof(ftype)));
+    cudaErrchk(cudaDeviceSynchronize());
+
+    cudaErrchk(cudaMemcpy(res.getData() + linearIdx, &val, sizeof(ftype), cudaMemcpyHostToDevice));
+    cudaErrchk(cudaDeviceSynchronize());
+  }
+}
diff --git a/src/backend/computational_graph/tensor_ops/cuda/tensor_ops_nodes.cuh b/src/backend/computational_graph/tensor_ops/cuda/tensor_ops_nodes.cuh
new file mode 100644
index 0000000..5568d2b
--- /dev/null
+++ b/src/backend/computational_graph/tensor_ops/cuda/tensor_ops_nodes.cuh
@@ -0,0 +1,24 @@
+/**
+ * @file tensor_ops_nodes.cuh
+ * @author Robert Baumgartner (r.baumgartner-1@tudelft.nl)
+ * @brief
+ * @version 0.1
+ * @date 2026-05-20
+ *
+ * @copyright Copyright (c) 2026
+ *
+ */
+
+#pragma once
+
+#ifndef __CUDA
+static_assert(false, "File should not be included without CUDA enabled");
+#endif // __CUDA
+
+#include "utility/global_params.h"
+#include "data_modeling/tensor.h"
+
+namespace cuda_impl {
+  void scalarMulBackward(Tensor& res, const Tensor& upstreamGrad, ftype factor);
+  void getterBackward(Tensor& res, ftype val, tensorSize_t linearIdx);
+}
diff --git a/src/backend/computational_graph/tensor_ops/matmul_node.cpp b/src/backend/computational_graph/tensor_ops/matmul_node.cpp
index 2237026..7edd0e7 100644
--- a/src/backend/computational_graph/tensor_ops/matmul_node.cpp
+++ b/src/backend/computational_graph/tensor_ops/matmul_node.cpp
@@ -16,6 +16,7 @@ using namespace cgraph;
 
 vector<shared_ptr<Tensor>> MatMulNode::backward(const Tensor& upstreamGrad) {
     assert(!upstreamGrad.getRequiresGrad());
+    // TODO: optimize operators
     return {
         make_shared<Tensor>(upstreamGrad.matmul(parents[1]->transpose(-2, -1))), 
         make_shared<Tensor>(parents[0]->transpose(-2, -1).matmul(upstreamGrad))
diff --git a/src/backend/computational_graph/tensor_ops/scalar_op_nodes.cpp b/src/backend/computational_graph/tensor_ops/scalar_op_nodes.cpp
index e0f52f0..fc70bfe 100644
--- a/src/backend/computational_graph/tensor_ops/scalar_op_nodes.cpp
+++ b/src/backend/computational_graph/tensor_ops/scalar_op_nodes.cpp
@@ -12,6 +12,11 @@
 #include "scalar_op_nodes.h"
 
 #include <utility>
+#include <stdexcept>
+
+#ifdef __CUDA
+#include "computational_graph/tensor_ops/cuda/tensor_ops_nodes.cuh"
+#endif
 
 using namespace std;
 using namespace cgraph;
@@ -25,8 +30,22 @@ vector<shared_ptr<Tensor>> cgraph::ScalarMulNode::backward(const Tensor& upstrea
   assert(!upstreamGrad.getRequiresGrad());
 
   auto res = make_shared<Tensor>(upstreamGrad.createDeepCopy());
-  for(tensorSize_t i=0; i<res->getSize(); i++){
-    res->set(res->get(i) * factor, i);
+  switch(res->getDevice()) {
+    case Device::CPU:
+    {
+      for(tensorSize_t i=0; i<res->getSize(); i++){
+        res->set(res->get(i) * factor, i);
+      }
+      break;
+    }
+    case Device::CUDA:
+    #ifdef __CUDA
+      cuda_impl::scalarMulBackward(*res, upstreamGrad, factor);
+    #else
+      __throw_invalid_argument("Not compiled with CUDA");
+    #endif
+      break;
   }
+
   return {std::move(res)};
 }
\ No newline at end of file
diff --git a/src/backend/data_modeling/cuda/tensor_ops.cu b/src/backend/data_modeling/cuda/tensor_ops.cu
new file mode 100644
index 0000000..d42b4d5
--- /dev/null
+++ b/src/backend/data_modeling/cuda/tensor_ops.cu
@@ -0,0 +1,275 @@
+/**
+ * @file tensor_ops.cu
+ * @author Robert Baumgartner (r.baumgartner-1@tudelft.nl)
+ * @brief 
+ * @version 0.1
+ * @date 2026-03-23
+ * 
+ * @copyright Copyright (c) 2026
+ * 
+ */
+
+#ifndef __CUDA
+static_assert(false, "File should not be compiled without CUDA enabled");
+#endif // __CUDA
+
+#include "data_modeling/tensor.h"
+
+#include "tensor_ops.cuh"
+#include "utility/cuda/cuda_common.cuh"
+#include "utility/utils.h"
+
+#include <thrust/fill.h>
+#include <thrust/device_ptr.h>
+
+namespace{
+  /**
+   * @brief Kernel for simple elementwise addition.
+   */
+  __global__ void elementwiseaddKernel(ftype* const res, const ftype* const left, const ftype* const right, const tensorSize_t size) {
+    int gid = blockDim.x * blockIdx.x + threadIdx.x;
+    if(gid>=size)
+      return;
+    
+    res[gid] = left[gid] + right[gid];
+  }
+
+  /**
+   * @brief Kernel for broadcasted addition, e.g. when adding a bias to a matrix.
+   */
+  __global__ void broadcastaddKernel(ftype* const res, const ftype* const matrix, const ftype* const vec, 
+                                     const tensorSize_t vectorSize, const tensorSize_t matrixSize) {
+    int gid = blockDim.x * blockIdx.x + threadIdx.x;
+    if(gid>=matrixSize)
+        return;
+
+    const int vectorIdx = gid % vectorSize;
+    res[gid] = matrix[gid] + vec[vectorIdx];
+  }
+
+  /**
+   * @brief Kernel for simple elementwise multiplication.
+   */
+  __global__ void elementwisemulKernel(ftype* const res, const ftype* const left, const ftype* const right, const tensorSize_t size) {
+    int gid = blockDim.x * blockIdx.x + threadIdx.x;
+    if(gid>=size)
+      return;
+    
+    res[gid] = left[gid] * right[gid];
+  }
+
+  /**
+   * @brief Kernel for scalar addition.
+   */
+  __global__ void scalaraddKernel(ftype* const res, const ftype* const left, ftype scalar, tensorSize_t size) {
+    int gid = blockDim.x * blockIdx.x + threadIdx.x;
+    if(gid>=size)
+      return;
+    
+    res[gid] = left[gid] + scalar;
+  }
+
+  /**
+   * @brief Kernel for scalar multiplication.
+   */
+  __global__ void scalarmulKernel(ftype* const res, const ftype* const left, ftype scalar, tensorSize_t size) {
+    int gid = blockDim.x * blockIdx.x + threadIdx.x;
+    if(gid>=size)
+      return;
+    
+    res[gid] = left[gid] * scalar;
+  }
+
+  /**
+   * @brief 2D matMul operation. Assumes that res is where result is written to, and res is zero-indexed.
+   * For higher dimensionalities the caller is responsible for proper offset computation!!! The same goes
+   * for left and right, which are also assumed to be 2D matrices, and offsets have to be computed by caller.
+   * 
+   * @param res The result matrix
+   * @param left Left matrix to be multiplied
+   * @param right Right matrix to multiply with
+   * @param leftRows n rows of 
+   * @param leftCols n columns of left matrix; leftCols == rightRows => rightRows not handed as param
+   * @param rightCols n columns of right matrix.
+   * @param resSize Size of the resulting 2D matrix.
+   */
+  __global__ void matMul2DKernel(ftype* const res, const ftype* const left, const ftype* const right,
+                               const tensorDim_t leftRows, const tensorDim_t leftCols,
+                               const tensorDim_t rightCols, const tensorSize_t resSize)
+  {
+    const int gid = blockDim.x * blockIdx.x + threadIdx.x;
+    if(gid >= resSize) 
+      return;
+    //const int tid = threadIdx.x;
+
+    // 48KB of shared mem -> ~12K floats of 4 bytes -> ~10 blocks per SM of shared memory is limit
+    //extern __shared__ ftype smem[];
+    //smem[tid] = 0.0f;
+    //__syncthreads();
+
+    const int resCol = gid % rightCols;
+    const int resRow = gid / rightCols;
+    const int leftBase = resRow * leftCols;
+
+    // C[i, j] = sum_{k=0}^{leftCols} A[i, k] * B[k, j]
+    ftype cij = 0;
+    for(int k = 0; k < leftCols; k++) {
+      //smem[tid] += left[leftBase + k] * right[k * rightCols + resCol];
+      cij += left[leftBase + k] * right[k * rightCols + resCol];
+    }
+
+    res[gid] = cij; // smem[tid];
+  }
+
+  /**
+   * @brief Strides in an outer size larger than one block. We use one thread per stride.
+   */
+  __global__ void sumOverDimsKernel(ftype* const res, const ftype* const input, tensorSize_t stride,
+                                    const tensorDim_t srcDimSize, const tensorSize_t size) {
+    const tensorSize_t gid = blockIdx.x * blockDim.x + threadIdx.x;
+    if(gid >= size) {
+      return;
+    } 
+
+    const tensorSize_t outerIdx = gid / stride;
+    const tensorSize_t batchOffset = outerIdx * stride * srcDimSize;
+
+    const int withinStrideIdx = gid % stride;
+
+    ftype sum = 0;
+    for(int k = 0; k < srcDimSize; k++) {
+      sum += input[batchOffset + k * stride + withinStrideIdx];
+    }
+
+    res[gid] = sum;
+  }
+
+  /**
+   * @brief Create a contiguous copy of src and copy it into dst. Used for reshaping, transposing, etc.
+   * 
+   * @param strides The original strides of src. With these strides src would be contiguous.
+   * @param contiguousStrides The new, contiguous strides of dst. We rearrange to match those in memory.
+   * @param ndim Number of dimension of both src and dst.
+   * @param size Total size of both src and dst.
+   */
+  __global__ void createContiguousCopyKernel(ftype* dst, const ftype* const src, const tensorSize_t* const strides,
+                                             const tensorDim_t* const dims, const int ndims, const tensorSize_t size)
+  {
+    const tensorSize_t flatIdx = blockIdx.x * blockDim.x + threadIdx.x;
+    if(flatIdx >= size) 
+      return;
+
+    tensorSize_t remainder = flatIdx;
+    tensorSize_t srcOffset = 0;
+    for (int i = ndims - 1; i >= 0; i--) {
+      tensorSize_t coord = remainder % dims[i];
+      srcOffset += coord * strides[i];
+      remainder /= dims[i];
+    }
+    dst[flatIdx] = src[srcOffset];
+  }
+}
+
+namespace cuda_impl {
+  void scalaradd(Tensor& res, const Tensor& src, ftype scalar) {
+    constexpr int threadsPerBlock = 256;
+    const int blocksPerGrid = (src.getSize() + threadsPerBlock - 1) / threadsPerBlock;
+
+    scalaraddKernel<<<blocksPerGrid, threadsPerBlock>>>(res.getData(), src.getData(), scalar, src.getSize());
+    cudaErrchk(cudaDeviceSynchronize());
+  }
+
+  void scalarmul(Tensor& res, const Tensor& src, ftype scalar) {
+    constexpr int threadsPerBlock = 256;
+    const int blocksPerGrid = (src.getSize() + threadsPerBlock - 1) / threadsPerBlock;
+
+    scalarmulKernel<<<blocksPerGrid, threadsPerBlock>>>(res.getData(), src.getData(), scalar, src.getSize());
+    cudaErrchk(cudaDeviceSynchronize());
+  }
+
+  void broadcastadd(Tensor& res, const Tensor& matrix, const Tensor& vec) {
+    const auto size = res.getSize();
+
+    constexpr int threadsPerBlock = 256;
+    const int blocksPerGrid = (size + threadsPerBlock - 1) / threadsPerBlock;
+
+    broadcastaddKernel<<<blocksPerGrid, threadsPerBlock>>>(
+      res.getData(), matrix.getData(), vec.getData(), vec.getDims()[0], matrix.getSize());
+    cudaErrchk(cudaDeviceSynchronize());
+  }
+
+  void elementwiseadd(Tensor& res, const Tensor& left, const Tensor& right) {
+    constexpr int threadsPerBlock = 256;
+    const int blocksPerGrid = (left.getSize() + threadsPerBlock - 1) / threadsPerBlock;
+
+    elementwiseaddKernel<<<blocksPerGrid, threadsPerBlock>>>(res.getData(), left.getData(), right.getData(), left.getSize());
+    cudaErrchk(cudaDeviceSynchronize());
+  }
+
+  void elementwisemul(Tensor& res, const Tensor& left, const Tensor& right) {
+    constexpr int threadsPerBlock = 256;
+    const int blocksPerGrid = (left.getSize() + threadsPerBlock - 1) / threadsPerBlock;
+
+    elementwisemulKernel<<<blocksPerGrid, threadsPerBlock>>>(res.getData(), left.getData(), right.getData(), left.getSize());
+    cudaErrchk(cudaDeviceSynchronize());
+  }
+
+  void matmul(Tensor& res, const Tensor& left, const Tensor& right) {
+    constexpr int threadsPerBlock = 256;
+    const int blocksPerGrid = (res.getSize() + threadsPerBlock - 1) / threadsPerBlock;
+    
+    // sizes of the 2D matrices respectively
+    const tensorSize_t leftSize = left.getDims().get(-1) * left.getDims().get(-2); 
+    const tensorSize_t rightSize = right.getDims().get(-1) * right.getDims().get(-2);
+    const tensorSize_t resSize = left.getDims().get(-2) * right.getDims().get(-1);
+
+    tensorSize_t leftOffset = 0;
+    tensorSize_t rightOffset = 0;
+    tensorSize_t resOffset = 0;
+
+    while(leftOffset < left.getSize()){
+      //const auto smemSize = min(resSize, threadsPerBlock) * sizeof(ftype);
+      //matMul2DKernel<<<blocksPerGrid, threadsPerBlock, smemSize>>>(res.getData() + resOffset, left.getData() + leftOffset, right.getData() + rightOffset, 
+      matMul2DKernel<<<blocksPerGrid, threadsPerBlock>>>(res.getData() + resOffset, left.getData() + leftOffset, right.getData() + rightOffset, 
+                                                         left.getDims().get(-2), left.getDims().get(-1), right.getDims().get(-1), resSize);
+
+      leftOffset += leftSize;
+      rightOffset += rightSize;
+      resOffset += resSize;
+    }
+    
+    cudaErrchk(cudaDeviceSynchronize());
+  }
+
+  void sumOverDims(Tensor& res, const Tensor& input, tensorDim_t dim) {
+    tensorSize_t stride = 1;
+    for(tensorDim_t i = dim + 1; i < input.getDims().nDims(); i++){
+      stride *= input.getDims()[i];
+    }
+
+    constexpr int threadsPerBlock = 256;
+    const int blocks = (res.getSize() + threadsPerBlock - 1) / threadsPerBlock;
+
+    sumOverDimsKernel<<<blocks, threadsPerBlock>>>(res.getData(), input.getData(), stride, input.getDims()[dim], res.getSize());
+    cudaErrchk(cudaDeviceSynchronize());
+  }
+
+  void scalarFill(Tensor& t, ftype value) {
+    ftype* ptr = t.getData();
+    thrust::fill(thrust::device_pointer_cast(ptr),
+                 thrust::device_pointer_cast(ptr + t.getSize()), value);
+    cudaErrchk(cudaDeviceSynchronize());
+  }
+
+  void createContiguousCopy(Tensor& res, const Tensor& src) {
+    assert(res.getSize()==src.getSize());
+
+    const auto size = src.getSize();
+    constexpr int threadsPerBlock = 256;
+    const int blocksPerGrid = (size + threadsPerBlock - 1) / threadsPerBlock;
+
+    createContiguousCopyKernel<<<blocksPerGrid, threadsPerBlock>>>(
+      res.getData(), src.getData(), src.getDims().getStrides().data(), src.getDims().data(), src.getDims().nDims(), size);
+    cudaErrchk(cudaDeviceSynchronize());
+  }
+}
diff --git a/src/backend/data_modeling/cuda/tensor_ops.cuh b/src/backend/data_modeling/cuda/tensor_ops.cuh
new file mode 100644
index 0000000..7634f46
--- /dev/null
+++ b/src/backend/data_modeling/cuda/tensor_ops.cuh
@@ -0,0 +1,39 @@
+/**
+ * @file tensor_ops.cuh
+ * @author Robert Baumgartner (r.baumgartner-1@tudelft.nl)
+ * @brief 
+ * @version 0.1
+ * @date 2026-03-23
+ * 
+ * @copyright Copyright (c) 2026
+ * 
+ */
+
+#pragma once
+
+#ifndef __CUDA
+static_assert(false, "File should not be included without CUDA enabled");
+#endif // __CUDA
+
+#include "utility/global_params.h"
+#include "data_modeling/dim_type.h"
+
+class Tensor;
+
+namespace cuda_impl {
+
+  // scalar ops
+  void scalaradd(Tensor& res, const Tensor& src, ftype scalar);
+  void scalarmul(Tensor& res, const Tensor& src, ftype scalar);
+
+  // matrix ops
+  void elementwiseadd(Tensor& res, const Tensor& left, const Tensor& right);
+  void broadcastadd(Tensor& res, const Tensor& matrix, const Tensor& vec);
+  void elementwisemul(Tensor& res, const Tensor& left, const Tensor& right);
+  void matmul(Tensor& res, const Tensor& left, const Tensor& right);
+  void sumOverDims(Tensor& res, const Tensor& input, tensorDim_t dim);
+
+  void scalarFill(Tensor& t, ftype value);
+
+  void createContiguousCopy(Tensor& res, const Tensor& src);
+}
diff --git a/src/backend/data_modeling/device.h b/src/backend/data_modeling/device.h
index 65f83a5..ea50669 100644
--- a/src/backend/data_modeling/device.h
+++ b/src/backend/data_modeling/device.h
@@ -11,9 +11,15 @@
 
 #pragma once
 
+#include <ostream>
+
 enum class Device {
     CPU,
     CUDA
 };
 
-const char* DeviceToString(Device d);
\ No newline at end of file
+const char* DeviceToString(Device d);
+
+inline std::ostream& operator<<(std::ostream& os, Device d) {
+  return os << DeviceToString(d);
+}
\ No newline at end of file
diff --git a/src/backend/data_modeling/dim_type.cpp b/src/backend/data_modeling/dim_type.cpp
index 684fde0..2eae45c 100644
--- a/src/backend/data_modeling/dim_type.cpp
+++ b/src/backend/data_modeling/dim_type.cpp
@@ -45,36 +45,55 @@ void Dimension::resize(const std::vector<tensorDim_t>& dims) {
 /**
  * @brief Swap along the two given dimensions.
  */
-void Dimension::swap(const tensorDim_t dim1, const tensorDim_t dim2) {
-  auto tmp = dims[dim1];
-  dims[dim1] = dims[dim2];
-  dims[dim2] = tmp;
+void Dimension::swap(int dim1, int dim2) {
+  if(dim1==dim2)
+    return;
+  
+  auto d1 = mapSignedIdx(dim1);
+  auto d2 = mapSignedIdx(dim2);
+  
+  std::swap(dims[d1], dims[d2]);
+  std::swap(strides[d1], strides[d2]);
 }
 
-Dimension::Dimension(const vector<tensorDim_t>& dims) : dims{dims} {
+Dimension::Dimension(const vector<tensorDim_t>& dims) 
+  : contiguousDims{dims}, contiguousStrides{makeStrides(dims)}, dims{dims}, strides{contiguousStrides} {
   size = multVector(dims);
+  lastDimIdx = dims.size()-1;
   assert(size>0);
 }
 
-Dimension::Dimension(const Dimension& other) : dims{other.dims}, size{other.size} { }
-
-Dimension& Dimension::operator=(const Dimension& other) {
-  if(this==&other) return *this;
+Dimension::Dimension(vector<tensorDim_t>&& dims, dim_t&& strides)
+  : contiguousDims{dims}, dims{contiguousDims}, contiguousStrides{strides}, strides{contiguousStrides} 
+{
+  size = multVector(dims);
+  lastDimIdx = dims.size()-1;
+  assert(size>0);
+}
 
-  dims = other.dims;
-  size = other.size;
+/**
+ * @brief Computes the strides as they are;
+ */
+Dimension::dim_t Dimension::makeStrides(const vector<tensorDim_t>& dims) const noexcept {
+  dim_t res;
+  const int lastDimIdx = dims.size() - 1;
 
-  return *this;
-}
+  tensorSize_t stride = 1;
+  res[lastDimIdx] = stride;
+  stride *= dims[lastDimIdx];
 
-Dimension::Dimension(Dimension&& other) noexcept : dims{move(other.dims)}, size{other.size} {}
+  for(int i = lastDimIdx - 1; i >= 0; i--){
+    res[i] = stride;
+    stride *= dims[i];
+  }
 
-Dimension& Dimension::operator=(Dimension&& other) noexcept {
-  if(this==&other) return *this;
+  return res;
+}
 
-  dims = move(other.dims);
-  size = other.size;
-  return *this;
+tensorSize_t Dimension::getStride(const int i) const noexcept {
+  if(i<0)
+    return strides[lastDimIdx + i + 1];
+  return strides[i];
 }
 
 /**
@@ -96,7 +115,16 @@ Dimension Dimension::collapseDimension(int idx) const {
   newDims.insert(newDims.end(), dims.begin(), dims.begin() + idx);
   newDims.insert(newDims.end(), dims.begin() + idx + 1, dims.end());
 
-  return Dimension(newDims);
+  dim_t newStrides{};
+  tensorDim_t strideIdx = 0;
+  for(tensorDim_t i=0; i<strides.size(); i++){
+    if(i==mappedIdx)
+      continue;
+
+    newStrides[strideIdx] = strides[i];
+  }
+
+  return Dimension(std::move(newDims), std::move(newStrides));
 }
 
 ostream& operator<<(ostream& os, const Dimension& d) noexcept {
diff --git a/src/backend/data_modeling/dim_type.h b/src/backend/data_modeling/dim_type.h
index 7b6bbed..8d5783f 100644
--- a/src/backend/data_modeling/dim_type.h
+++ b/src/backend/data_modeling/dim_type.h
@@ -14,34 +14,56 @@
 #include "utility/global_params.h"
 
 #include <vector>
+#include <array>
 
 #include <iostream>
 #include <cassert>
 
 class Dimension final {
+  using dim_t = std::array<tensorSize_t, MAX_NDIMS>;
+
   private:
+    // those two indicate the structure of the contiguous data that lies underneath
+    // WARNING: Should only be set in ctor. We did not make them constant, so that we can use copy c'tor
+    std::vector<tensorDim_t> contiguousDims;
+    dim_t contiguousStrides;
+
     std::vector<tensorDim_t> dims;
-    tensorSize_t size = 0;
+    dim_t strides;
+
+    tensorDim_t lastDimIdx; // look up end in strides/dims
+    tensorSize_t size = 0; // total size of tensor
 
+    dim_t makeStrides(const std::vector<tensorDim_t>& dims) const noexcept;
     tensorSize_t multVector(const std::vector<tensorDim_t>& dims) const noexcept;
+    
+    Dimension(std::vector<tensorDim_t>&& dims, dim_t&& strides);
+
+    tensorDim_t mapSignedIdx(int idx) const {
+      if(idx < 0) {
+        // -1 is last idx, -2 second last and so forth
+        return lastDimIdx + idx + 1;
+      }
+      return idx;
+    }
 
   public:
-    /**
-     * @brief Explicit default ctor, so that dims is zero initialized.
-     * Otherwise we will encounter undefined behavior.
-     */
     Dimension(const std::vector<tensorDim_t>& dims);
 
-    Dimension(const Dimension& other);
-    Dimension& operator=(const Dimension& other);
+    Dimension(const Dimension& other) = default;
+    Dimension& operator=(const Dimension& other) = default;
 
-    Dimension(Dimension&& other) noexcept;
-    Dimension& operator=(Dimension&& other) noexcept;
+    Dimension(Dimension&& other) noexcept = default;
+    Dimension& operator=(Dimension&& other) noexcept = default;
 
     ~Dimension() noexcept = default;
 
     Dimension collapseDimension(int idx) const;
 
+    bool inOriginalState() const noexcept { 
+      return contiguousDims == dims && contiguousStrides == strides; 
+    }
+
     void resize(const std::vector<tensorDim_t>& dims);
       
     tensorSize_t getSize() const noexcept {
@@ -54,18 +76,20 @@ class Dimension final {
 
     tensorDim_t operator[](int idx) const {
       assert(size>0);
-      if(idx<0){
-        idx = dims.size() + idx; // -1 is last idx, -2 second last and so forth
-      }
-
-      return dims[idx];
+      return dims[mapSignedIdx(idx)];
     }
 
+    const tensorDim_t* data() const noexcept { return dims.data(); }
+    const auto getStrides() const noexcept { return strides; }
+    const auto getContiguousStrides() const noexcept { return contiguousStrides; }
+
+    tensorSize_t getStride(int i) const noexcept;
+
     std::vector<tensorDim_t> toVector() const noexcept{
       return dims;
     }
 
-    void swap(const tensorDim_t dim1, const tensorDim_t dim2);
+    void swap(int dim1, int dim2);
 
     size_t nDims() const noexcept {
       return dims.size();
diff --git a/src/backend/data_modeling/tensor.cpp b/src/backend/data_modeling/tensor.cpp
index e398c2c..755d939 100644
--- a/src/backend/data_modeling/tensor.cpp
+++ b/src/backend/data_modeling/tensor.cpp
@@ -14,10 +14,19 @@
 #include "computational_graph/graph_node.h"
 #include "computational_graph/topological_sort.h"
 
+#include "utility/utils.h"
+
 #include <utility>
 #include <limits>
 #include <cstring>
 
+#include <format>
+
+#ifdef __CUDA
+#include "utility/cuda/cuda_common.cuh"
+#include "data_modeling/cuda/tensor_ops.cuh"
+#endif
+
 using namespace std;
 
 /******************************************************************** 
@@ -51,11 +60,32 @@ Tensor::tensorValues_t::~tensorValues_t() noexcept {
       free(values);
       break;
     case Device::CUDA:
-      std::__throw_invalid_argument("Cuda destructor not implemented yet.");
+      #ifdef __CUDA
+        cudaErrchk(cudaFree(values));
+      #else
+        std::__throw_invalid_argument("Not compiled with CUDA.");
+      #endif
+      break;
+  }
+}
+
+void Tensor::tensorValues_t::resize(const tensorSize_t size) {
+  this->size = size;
+  switch (device) {
+    case Device::CPU:
+      values = static_cast<ftype*>(std::malloc(this->size * sizeof(ftype)));
+      break;
+    case Device::CUDA:
+      #ifdef __CUDA
+        cudaErrchk(cudaMalloc((void**) &values, this->size * sizeof(ftype)));
+      #else
+        std::__throw_invalid_argument("Not compiled with CUDA.");
+      #endif
       break;
   }
 }
 
+
 /**
  * @brief Copy from pointer into this object.
  */
@@ -66,7 +96,12 @@ void Tensor::tensorValues_t::copyFromRaw(const ftype* src, tensorSize_t n) {
       std::memcpy(values, src, n * sizeof(ftype));
       break;
     case Device::CUDA:
-      __throw_runtime_error("copyFromRaw not implemented for CUDA");
+      #ifdef __CUDA
+        cudaErrchk(cudaMemcpy(values, src, n*sizeof(ftype), cudaMemcpyDeviceToDevice));
+      #else
+        __throw_runtime_error("Not compiled with CUDA");
+      #endif
+      break;
   }
 }
 
@@ -80,12 +115,14 @@ void Tensor::tensorValues_t::copyValues(Tensor::tensorValues_t& target) const {
 
   switch(device){
     case Device::CPU:
-      for(tensorSize_t i=0; i<size; i++){
-        target[i] = values[i];
-      }
+      std::memcpy(target.values, values, size * sizeof(ftype));
       break;
     case Device::CUDA:
-      __throw_runtime_error("Deep copy not implemented for CUDA");
+      #ifdef __CUDA
+        cudaErrchk(cudaMemcpy(target.values, values, size * sizeof(ftype), cudaMemcpyDeviceToDevice));
+      #else
+        __throw_runtime_error("Not compiled with CUDA");
+      #endif
       break;
   }
 }
@@ -96,15 +133,22 @@ void Tensor::tensorValues_t::copyValues(Tensor::tensorValues_t& target) const {
 void Tensor::tensorValues_t::copyValues(tensorValues_t& target, tensorSize_t low, 
                                         tensorSize_t high, tensorSize_t targetOffset) const {
   assert(target.size >= high - low);
+  if(high<low){
+    __throw_invalid_argument(
+      std::format("high argument should be higher than low, but received {0} and {1}", high, low).c_str()
+    );
+  }
 
   switch(device){
     case Device::CPU:
-      for(tensorSize_t i=0; i<high-low; i++){
-        target[targetOffset+i] = values[low+i];
-      }
+      std::memcpy(target.values+targetOffset, values+low, (high-low) * sizeof(ftype));
       break;
     case Device::CUDA:
-      __throw_runtime_error("CUDA not implemented for slicing");
+      #ifdef __CUDA
+        cudaErrchk(cudaMemcpy(target.values+targetOffset, values+low, (high-low) * sizeof(ftype), cudaMemcpyDeviceToDevice));
+      #else
+        __throw_runtime_error("Not compiled with CUDA");
+      #endif
       break;
   }
 }
@@ -130,16 +174,61 @@ void Tensor::tensorValues_t::copyValues(tensorValues_t& target, span<const tenso
       break; 
     }
     case Device::CUDA:
-      __throw_runtime_error("CUDA not implemented for slicing");
+      #ifdef __CUDA
+        // TODO: we can do streams here, that's why synchronize
+        tensorSize_t targetOffset = 0;
+        for(tensorDim_t idx: indices){
+          tensorSize_t thisOffset = idx * sizeOfDim;
+          copyValues(target, thisOffset, thisOffset+sizeOfDim, targetOffset);
+          targetOffset += sizeOfDim;
+        }
+        cudaErrchk(cudaDeviceSynchronize());
+      #else 
+        __throw_invalid_argument("Not compiled with CUDA");
+      #endif
       break;
   }
 }
 
 void Tensor::tensorValues_t::setDevice(const Device d) noexcept {
-  __throw_runtime_error("Not implemented yet.");
-  device = d;
+  #ifndef __CUDA
+    cerr << "setDevice called when CUDA not enabled. Doing nothing." << endl;
+    return;
+  #else 
+    if(d==device){
+      return;
+    }
+
+    switch(device){
+      case Device::CPU:
+        if(d == Device::CUDA) {
+          ftype* tmp;
+          cudaErrchk(cudaMalloc((void**) &tmp, size * sizeof(ftype)));
+          cudaErrchk(cudaMemcpy(tmp, values, size * sizeof(ftype), cudaMemcpyHostToDevice));
+          free(values);
+          values = tmp;
+        }
+        else {
+          __throw_runtime_error((string("Requested unexpected device ") + DeviceToString(d)).c_str());
+        }
+        break;
+      case Device::CUDA:
+        if(d == Device::CPU) {
+          ftype* tmp = static_cast<ftype*>(std::malloc(size * sizeof(ftype)));
+          cudaErrchk(cudaMemcpy(tmp, values, size * sizeof(ftype), cudaMemcpyDeviceToHost));
+          cudaErrchk(cudaFree(values));
+          values = tmp;
+        }
+        else {
+          __throw_runtime_error((string("Requested unexpected device ") + DeviceToString(d)).c_str());
+        }
+        break;
+    }
+    device = d;
+  #endif
 }
 
+
 Device Tensor::tensorValues_t::getDevice() const noexcept {
   return device;
 }
@@ -160,46 +249,29 @@ Tensor::tensorValues_t::operator bool() const noexcept {
   return values != nullptr;
 }
 
-void Tensor::tensorValues_t::addOtherCpu(const Tensor::tensorValues_t& other) noexcept {
-  for(tensorSize_t i=0; i<this->size; i++){
-    this->values[i] += other.values[i];
-  }
-}
-
-Tensor::tensorValues_t& 
-Tensor::tensorValues_t::operator+=(const Tensor::tensorValues_t& other) {
-  assert(this->size==other.size && this->device == other.device);
-  
-  switch(device) {
-    case Device::CPU:
-      addOtherCpu(other);
-      break;
-    case Device::CUDA:
-      __throw_invalid_argument("CUDA not supported yet for += operation");
-      break;
-  }
-
-  return *this;
-}
-
-ftype& Tensor::tensorValues_t::operator[](const tensorSize_t idx) {
-  if(idx >= size)
-    throw std::out_of_range("Out of range for tensor");
-
-  if(device!=Device::CPU){
-    __throw_invalid_argument("Operator [] only implemented for CPU");
-  }
-  return values[idx];
-}
 
 ftype Tensor::tensorValues_t::operator[](const tensorSize_t idx) const {
   if(idx >= size)
     throw std::out_of_range("Out of range for tensor");
 
-  if(device!=Device::CPU){
-    __throw_invalid_argument("Operator [] only implemented for CPU");
+  switch(device){
+    case Device::CPU:
+      return values[idx];
+    case Device::CUDA:
+      #ifdef __CUDA
+      {
+        ftype res;
+        cudaErrchk(cudaMemcpy(&res, values + idx, sizeof(ftype), cudaMemcpyDeviceToHost));
+        return res;
+      }
+      #else
+        __throw_invalid_argument("Not compiled with CUDA");
+        break;
+      #endif
   }
-  return values[idx];
+
+  __throw_runtime_error("Should never reach here.");
+  return 0; // suppress warnings
 }
 
 void Tensor::tensorValues_t::set(ftype v, tensorSize_t idx) {
@@ -209,23 +281,35 @@ void Tensor::tensorValues_t::set(ftype v, tensorSize_t idx) {
   switch(device){
     case Device::CPU:
       values[idx] = v;
-      return;
+      break;
     case Device::CUDA:
-      __throw_runtime_error("Not implemented for CUDA yet");
+      #ifdef __CUDA
+        cudaErrchk(cudaMemcpy(values + idx, &v, sizeof(ftype), cudaMemcpyHostToDevice));
+      #else
+        __throw_invalid_argument("Not compiled with CUDA");
+      #endif
+      break;
   }
-
-  __throw_runtime_error("Should never reach here.");
 }
 
 ftype Tensor::tensorValues_t::get(tensorSize_t idx) {
   if(idx >= size)
     throw std::out_of_range("Out of range for tensor");
-  
+
   switch(device){
     case Device::CPU:
       return values[idx];
     case Device::CUDA:
-      __throw_runtime_error("Not implemented for CUDA yet");
+      #ifdef __CUDA
+      {
+        ftype res;
+        cudaErrchk(cudaMemcpy(&res, values + idx, sizeof(ftype), cudaMemcpyDeviceToHost));
+        return res;
+      }
+      #else
+        __throw_invalid_argument("Not compiled with CUDA");
+        break;
+      #endif
   }
 
   __throw_runtime_error("Should never reach here.");
@@ -257,6 +341,17 @@ Tensor& Tensor::operator=(Tensor&& other) noexcept {
   return *this;
 }
 
+/**
+ * @brief Only createShallowCopy is supposed to access this ctor.
+ * Like a copy ctor, shares recourses (grad, cgNode, values) with 
+ * original (other) tensor.
+ */
+Tensor::Tensor(const Tensor& other, [[maybe_unused]] shallowCopyToken)
+  : dims(other.dims), cgNode{other.cgNode}, values{other.values}, 
+    grads{other.grads}, requiresGrad{other.requiresGrad} 
+{
+}
+
 /**
  * @brief Creates an empty copy of this tensor.
  * Metadata all filled, but gradients not initialized, and 
@@ -266,6 +361,72 @@ Tensor Tensor::createEmptyCopy() const {
   auto res = Tensor(dims, values->getDevice(), requiresGrad);
   return res;
 }
+
+/**
+ * @brief Creates a shallow copy. New tensor object, but all resources 
+ * (grads, cgNode, values) are shared with this tensor. Useful for optimization.
+ */
+Tensor Tensor::createShallowCopy() const {
+  auto res = Tensor(*this, shallowCopyToken{});
+  return res;
+}
+
+/**
+ * @brief Creates a linearized copy of this tensor. New memory allocations for 
+ * values, and values will be copied so that the memory layout reflects the 
+ * current shape of the tensor.
+ */
+Tensor Tensor::createContiguousCopy() const {
+  auto res = createEmptyCopy();
+  
+  switch(values->getDevice()) {
+    case Device::CPU:
+    {
+      for(tensorSize_t flatIdx = 0; flatIdx < values->getSize(); flatIdx++) {
+        tensorSize_t remainder = flatIdx;
+        tensorSize_t srcOffset = 0;
+
+        for(int i = dims.nDims() - 1; i >= 0; i--) {
+          tensorSize_t coord = remainder % dims[i];
+          srcOffset += coord * dims.getStride(i);
+          remainder /= dims[i];
+        }
+
+        res.values->data()[flatIdx] = (*values)[srcOffset];
+      }
+      break;
+    }
+    case Device::CUDA:
+    {
+      #ifdef __CUDA
+        cuda_impl::createContiguousCopy(res, *this);
+      #else
+        __throw_runtime_error("Not compiled with CUDA");
+      #endif
+      break;
+    }
+  }
+
+  return res;
+}
+
+void Tensor::makeContiguous() const {
+  if(isContiguous())
+    return;
+
+  auto tmp = createContiguousCopy();
+  values = std::move(tmp.values);
+  dims = std::move(tmp.dims);
+
+  // TODO: perhaps we do this one lazily too?
+  if(grads){
+    if(grads->hasGrads())
+      __throw_runtime_error("Grads should never have grads themselves");
+    
+    grads->makeContiguous();
+  }
+}
+
 /**
  * @brief Does a deep copy, but omits gradient and computational graph information.
  */
@@ -275,26 +436,19 @@ Tensor Tensor::createDeepCopy() const {
   auto res = Tensor(dims, values->getDevice(), requiresGrad);
   values->copyValues(*res.values);
 
-  /* if(grads){
-    res.grads = make_shared<Tensor>( grads->createDeepCopy() ); // TODO: do we want this?
-  } */
-
   assert(!res.grads); // TODO: check if this makes sense
   assert(!res.cgNode); // TODO: do we want to give it pointer of same node?
   return res;
 }
 
-
 /**
- * @brief Scalar multiplication, capable of broadcasting. First argument assumed
- * to be the scalar tensor.
+ * @brief For convenience. Needed for other classes to perform their CUDA operations.
+ * Avoids moving all CUDA code into tensor.
  */
-Tensor Tensor::multiplyScalar(const Tensor& scalar, const Tensor& right) noexcept {
-  Tensor res(right.dims, right.values->getDevice(), false);
-  for(int i=0; i<right.getSize(); ++i){
-    (*res.values)[i] = (*scalar.values)[0] * (*right.values)[i];
-  }
-  return res;
+ftype* Tensor::getData() const noexcept {
+  if(values->getDevice() == Device::CPU)
+    __throw_runtime_error("Should only be called on CUDA tensor.");
+  return values->data();
 }
 
 /**
@@ -315,26 +469,39 @@ Tensor Tensor::matMulImpl(const Tensor& left, const Tensor& right) {
 
   // broadcasting
   auto resDims = left.dims.nDims() > right.dims.nDims() ? left.dims.toVector() : right.dims.toVector();
-  resDims[resDims.size()-2] = left.dims.get(-2); // rows
-  resDims[resDims.size()-1] = right.dims.get(-1); // cols
+  resDims[resDims.size() - 2] = left.dims.get(-2); // rows
+  resDims[resDims.size() - 1] = right.dims.get(-1); // cols
 
   Tensor res(resDims, left.values->getDevice(), false);
 
-  // sizes of the 2D matrices respectively
-  const tensorSize_t leftSize = left.dims.get(-1) * left.dims.get(-2); 
-  const tensorSize_t rightSize = right.dims.get(-1) * right.dims.get(-2);
-  const tensorSize_t resSize = left.dims.get(-2) * right.dims.get(-1);
-
-  tensorSize_t leftOffset = 0;
-  tensorSize_t rightOffset = 0;
-  tensorSize_t resOffset = 0;
-
-  while(leftOffset < left.getSize()){
-    matMul2DCpu(res, left, right, resOffset, leftOffset, rightOffset);
+  switch(left.values->getDevice()){
+    case Device::CPU:
+    {
+      // sizes of the 2D matrices respectively
+      const tensorSize_t leftSize = left.dims.get(-1) * left.dims.get(-2); 
+      const tensorSize_t rightSize = right.dims.get(-1) * right.dims.get(-2);
+      const tensorSize_t resSize = left.dims.get(-2) * right.dims.get(-1);
+
+      tensorSize_t leftOffset = 0;
+      tensorSize_t rightOffset = 0;
+      tensorSize_t resOffset = 0;
+
+      while(leftOffset < left.getSize()){
+        matMul2DCpu(res, left, right, resOffset, leftOffset, rightOffset);
+        leftOffset += leftSize;
+        rightOffset += rightSize;
+        resOffset += resSize;
+      }
 
-    leftOffset += leftSize;
-    rightOffset += rightSize;
-    resOffset += resSize;
+      break;
+    }
+    case Device::CUDA:
+      #ifdef __CUDA
+        cuda_impl::matmul(res, left, right);
+      #else
+        __throw_invalid_argument("Not compiled with CUDA");
+      #endif
+      break;
   }
 
   return res;
@@ -351,22 +518,23 @@ void Tensor::matMul2DCpu(Tensor& res, const Tensor& left, const Tensor& right, c
   const auto nRowsRight = static_cast<tensorSize_t>(right.dims.get(-2));
   const auto nColsRight = static_cast<tensorSize_t>(right.dims.get(-1));
 
-  tensorSize_t resIdx = resOffset;
-  for(tensorSize_t lrow=0; lrow<nRowsLeft; lrow++){
+  for(tensorSize_t lrow = 0; lrow < nRowsLeft; lrow++){
+    const tensorSize_t resRowOffset = resOffset + lrow * nColsRight;
+    const tensorSize_t leftRowOffset = leftOffset + lrow * nColsLeft;
 
-    for(tensorSize_t rcol=0; rcol<nColsRight; rcol++){
-      tensorSize_t leftIdx = leftOffset + lrow * nColsLeft;
-      tensorSize_t rightIdx = rightOffset + rcol;
+    tensorSize_t rightIdx = rightOffset;
+    // res likely has undefined memory content
+    for(tensorSize_t rrow = 0; rrow < nColsRight; rrow++){
+      res.values->data()[resRowOffset + rrow] = left.values->data()[leftRowOffset] * right.values->data()[rightIdx];
+      rightIdx++;
+    }
 
-      ftype scalar = 0.0;
-      for(tensorSize_t lcol=0; lcol<nColsLeft; lcol++){
-        scalar += (*left.values)[leftIdx] * (*right.values)[rightIdx];
-        leftIdx++;
-        rightIdx += nColsRight;
+    for(tensorSize_t lcol=1; lcol<nColsLeft; lcol++){
+      const auto leftIdx = leftRowOffset + lcol;
+      for(tensorSize_t rrow = 0; rrow < nColsRight; rrow++){
+        res.values->data()[resRowOffset + rrow] += left.values->data()[leftIdx] * right.values->data()[rightIdx];
+        rightIdx++;
       }
-      
-      (*res.values)[resIdx] = scalar;
-      resIdx++;
     }
   }
 }
@@ -375,16 +543,11 @@ void Tensor::matMul2DCpu(Tensor& res, const Tensor& left, const Tensor& right, c
  * @brief Matrix multiplication.
  */
 Tensor Tensor::matmul(const Tensor& other) const {
-  assert(values->getDevice()==other.values->getDevice());
-  if(values->getDevice()==Device::CUDA){
-    __throw_invalid_argument("Multiplication not implemented on CUDA");
-  }
-
   if(values->getDevice()!=other.values->getDevice()){
     __throw_runtime_error("Tensors on different devices.");
   }
-
-  return matMulImpl(*this, other);
+  
+  return matMulImpl(getContiguous(), other.getContiguous());
 }
 
 /**
@@ -395,10 +558,6 @@ Tensor Tensor::matmul(const Tensor& other) const {
  * other.dims == (dimN) && this->dims == (dim0, dim1,..., dimN).
  */
 Tensor Tensor::operator+(const Tensor& other) const {
-  if(values->getDevice()==Device::CUDA){
-    __throw_invalid_argument("Addition not implemented on CUDA");
-  }
-
   if(this->dims != other.dims && 
     !(other.dims.nDims() == 1 && other.dims.get(0) == dims.get(-1))){
     __throw_invalid_argument("Tensors need matching dimensions");
@@ -407,24 +566,42 @@ Tensor Tensor::operator+(const Tensor& other) const {
     __throw_runtime_error("Tensors on different devices.");
   }
 
+  makeContiguous();
+  other.makeContiguous();
   Tensor res(dims, values->getDevice());
 
-  if(dims==other.dims){
-    // elementwise add
-    for(tensorSize_t i=0; i<values->getSize(); i++){
-      (*res.values)[i] = (*values)[i] + (*other.values)[i];
-    }
-  }
-  else { [[likely]]
-    // broadcasted add
-    const auto stride = static_cast<tensorSize_t>(other.dims.get(0));
-    for(tensorSize_t offset=0; offset<values->getSize(); offset+=stride){
-      for(tensorSize_t i=0; i<stride; i++){
-        (*res.values)[offset+i] = (*values)[offset+i] + (*other.values)[i];
+  switch(values->getDevice()){
+    case Device::CPU:
+      if(dims==other.dims) [[unlikely]] {
+        // elementwise add
+        for(tensorSize_t i = 0; i < values->getSize(); i++){
+          res.values->data()[i] = values->data()[i] + other.values->data()[i];
+        }
+        return res;
       }
-    }
+      else [[likely]] {
+        // broadcasted add
+        const auto stride = static_cast<tensorSize_t>(other.dims.get(0)); // other is a vector
+        for(tensorSize_t offset=0; offset<values->getSize(); offset += stride){
+          for(tensorSize_t i = 0; i < stride; i++){
+            res.values->data()[offset + i] = values->data()[offset + i] + other.values->data()[i];
+          }
+        }
+      }
+      break;
+    case Device::CUDA:
+      #ifdef __CUDA
+        if(dims==other.dims) [[unlikely]] {
+          cuda_impl::elementwiseadd(res, *this, other);
+        }
+        else [[likely]] {
+          cuda_impl::broadcastadd(res, *this, other);
+        }
+      #else
+        __throw_runtime_error("Not compiled with CUDA");
+      #endif
+      break;
   }
-
   return res;
 }
 
@@ -439,19 +616,6 @@ Tensor Tensor::add(const Tensor& other) const {
  * @brief Elementwise multiplication.
  */
 Tensor Tensor::operator*(const Tensor& other) const {
-  assert(values->getDevice()==other.values->getDevice());
-  if(values->getDevice()==Device::CUDA){
-    __throw_invalid_argument("Multiplication not implemented on CUDA");
-  }
-
-  // TODO: check what to do about these two gradients and if you want broadcasting here at all
-/*   if(other.dims.getSize()==1){
-    return multiplyScalar(other, *this);
-  }
-  else if(dims.getSize()==1){
-    return multiplyScalar(*this, other);
-  } */
-
   if(this->dims != other.dims){
     __throw_invalid_argument("Tensors need same dimensions");
   }
@@ -459,10 +623,23 @@ Tensor Tensor::operator*(const Tensor& other) const {
     __throw_runtime_error("Tensors on different devices.");
   }
 
-  assert(values->getSize()==other.values->getSize());
-  Tensor res(dims, values->getDevice(), false);
-  for(tensorSize_t i=0; i<values->getSize(); i++){
-    (*res.values)[i] = (*values)[i] * (*other.values)[i];
+  makeContiguous();
+  other.makeContiguous();
+  Tensor res(dims, values->getDevice(), requiresGrad);
+
+  switch(values->getDevice()){
+    case Device::CPU:
+      for(tensorSize_t i = 0; i < values->getSize(); i++){
+        res.values->data()[i] = values->data()[i] * other.values->data()[i];
+      }
+      break;
+    case Device::CUDA:
+      #ifdef __CUDA
+        cuda_impl::elementwisemul(res, *this, other);
+      #else
+        __throw_runtime_error("Not compiled with CUDA");
+      #endif
+      break;
   }
 
   return res;
@@ -475,51 +652,114 @@ Tensor Tensor::elementwiseMul(const Tensor& other) const {
   return *this * other;
 }
 
-Tensor Tensor::operator*(ftype scalar) const {
-  Tensor res(dims, values->getDevice(), false);
-  for (tensorSize_t i = 0; i < values->getSize(); ++i) {
-    (*res.values)[i] = (*values)[i] * scalar;
+Tensor& Tensor::operator+=(const Tensor& other) {
+  assert(dims == other.dims && values->getDevice() == other.values->getDevice());
+
+  switch(values->getDevice()) {
+    case Device::CPU:
+      for(tensorSize_t i = 0; i < values->getSize(); i++)
+        values->data()[i] += other.values->data()[i];
+      break;
+    case Device::CUDA:
+      #ifdef __CUDA
+        cuda_impl::elementwiseadd(*this, *this, other);
+      #else
+        __throw_runtime_error("Not compiled with CUDA");
+      #endif
+      break;
+  }
+  return *this;
+}
+
+Tensor Tensor::operator*(const ftype scalar) const {
+  Tensor res(dims, values->getDevice(), requiresGrad);
+  switch(values->getDevice()){
+    case Device::CPU:
+      for (tensorSize_t i = 0; i < values->getSize(); ++i) {
+        res.values->data()[i] = values->data()[i] * scalar;
+      }
+      break;
+    case Device::CUDA:
+      #ifdef __CUDA
+        cuda_impl::scalarmul(res, *this, scalar);
+      #else
+        __throw_runtime_error("Not compiled with CUDA");
+      #endif
+      break;
   }
 
   return res;
 }
 
-Tensor Tensor::operator/(ftype scalar) const {
+Tensor Tensor::operator/(const ftype scalar) const {
   if(scalar==0.0){
     __throw_runtime_error("Cannot divide by zero.");
   }
 
-  Tensor res(dims, values->getDevice(), false);
-  for (tensorSize_t i = 0; i < values->getSize(); ++i) {
-    (*res.values)[i] = (*values)[i] / scalar;
+  Tensor res(dims, values->getDevice(), requiresGrad);
+  switch(values->getDevice()){
+    case Device::CPU:
+      for (tensorSize_t i = 0; i < values->getSize(); ++i) {
+        res.values->data()[i] = values->data()[i] / scalar;
+      }
+      break;
+    case Device::CUDA:
+      #ifdef __CUDA
+        cuda_impl::scalarmul(res, *this, 1 / scalar);
+      #else 
+        __throw_runtime_error("Not compiled with CUDA");
+      #endif
+      break;
   }
 
   return res;
 }
 
-Tensor Tensor::operator+(ftype scalar) const {
-  Tensor res(dims, values->getDevice(), false);
-  for (tensorSize_t i = 0; i < values->getSize(); ++i) {
-    (*res.values)[i] = (*values)[i] + scalar;
+Tensor Tensor::operator+(const ftype scalar) const {
+  Tensor res(dims, values->getDevice(), requiresGrad);
+  switch(values->getDevice()){
+    case Device::CPU:
+      for (tensorSize_t i = 0; i < values->getSize(); ++i) {
+        res.values->data()[i] = values->data()[i] + scalar;
+      }
+      break;
+    case Device::CUDA:
+      #ifdef __CUDA
+        cuda_impl::scalaradd(res, *this, scalar);
+      #else
+        __throw_runtime_error("Not compiled with CUDA");
+      #endif
+      break;
   }
 
   return res;
 }
 
-Tensor Tensor::operator-(ftype scalar) const {
-  Tensor res(dims, values->getDevice(), false);
-  for (tensorSize_t i = 0; i < values->getSize(); ++i) {
-    (*res.values)[i] = (*values)[i] - scalar;
+Tensor Tensor::operator-(const ftype scalar) const {
+  Tensor res(dims, values->getDevice(), requiresGrad);
+  switch(values->getDevice()){
+    case Device::CPU:
+      for (tensorSize_t i = 0; i < values->getSize(); ++i) {
+        res.values->data()[i] = values->data()[i] - scalar;
+      }
+      break;
+    case Device::CUDA:
+      #ifdef __CUDA
+        cuda_impl::scalaradd(res, *this, -scalar);
+      #else 
+        __throw_runtime_error("Not compiled with CUDA");
+      #endif
+      break;
   }
 
   return res;
 }
 
-Tensor operator*(ftype scalar, const Tensor& tensor) {
+Tensor operator*(const ftype scalar, const Tensor& tensor) {
   return tensor * scalar;
 }
 
-Tensor operator+(ftype scalar, const Tensor& tensor) {
+Tensor operator+(const ftype scalar, const Tensor& tensor) {
   return tensor + scalar;
 }
 
@@ -531,24 +771,23 @@ void Tensor::backward() {
     __throw_runtime_error("Invoking backward on Tensor not created by a differentiable operation");
   }
 
-  // check this one out for sure
+  // last node has no incoming gradients -> factor 1
   if (!grads) {
-    grads = make_unique<Tensor>(dims, values->getDevice(), false);
-    for(tensorSize_t i=0; i<values->getSize(); i++){
-      (*grads->values)[i] = 1;
-    }
+    grads = make_shared<Tensor>(dims, values->getDevice(), false);
+    grads->reset(1);
   }
 
   vector<Tensor*> sortedTensors = cgraph::TopologicalSort::reverseSort(this);
-  for(auto tPtr: sortedTensors){
+  for(auto tPtr : sortedTensors){
     auto& tensor = *tPtr;
+    tensor.makeContiguous();
     assert(tensor.grads && !tensor.grads->requiresGrad); // gradient should not require grad
 
     auto incomingGrads = tensor.cgNode->backward(*tensor.grads);
 
     const auto& parents = tensor.cgNode->getParents();
 
-    for(size_t i=0; i<parents.size(); i++){
+    for(size_t i = 0; i < parents.size(); i++){
       auto parent = parents[i];
       if(!parent->requiresGrad){
         continue;
@@ -557,7 +796,7 @@ void Tensor::backward() {
         parent->grads = incomingGrads[i];
       }
       else{
-        *parent->grads->values += *incomingGrads[i]->values;
+        *parent->grads += *incomingGrads[i];
       }
     }
   }
@@ -575,7 +814,7 @@ shared_ptr<Tensor> Tensor::getGrads() const {
  * NumPy we map from the end to the beginning in that case. 
  */
 tensorDim_t Tensor::mapDim(const int dim, const Dimension& dims) {
-  if(dim>=0){
+  if(dim >= 0){
     return dim;
   }
   else if(dim + dims.nDims() < 0){
@@ -586,172 +825,22 @@ tensorDim_t Tensor::mapDim(const int dim, const Dimension& dims) {
   return dims.nDims() + dim;
 }
 
-void Tensor::transposeImpl(Tensor& target, const int dim1, const int dim2) const noexcept {
-    assert(values->getSize() == target.values->getSize() && dims.nDims()==target.dims.nDims());
-    if(target.getDevice() == Device::CUDA) {
-        __throw_runtime_error("Transposition for CUDA not implemented yet");
-    }
-    if(dim1 == dim2) {
-        return;
-    }
-
-    const auto& source = *this; // easier to read
-    
-    auto dim1Mapped = mapDim(dim1, source.dims);
-    auto dim2Mapped = mapDim(dim2, source.dims);
-    
-    const int numDims = source.dims.nDims();
-    std::vector<tensorSize_t> indices(numDims, 0);
-    std::vector<tensorSize_t> dimSizes(numDims);
-    std::vector<tensorSize_t> sourceStrides(numDims);
-    
-    // strides for source
-    tensorSize_t stride = 1;
-    for(int d = numDims - 1; d >= 0; d--) {
-        dimSizes[d] = source.dims.get(d);
-        sourceStrides[d] = stride;
-        stride *= dimSizes[d];
-    }
-    
-    // strides for target
-    std::vector<tensorSize_t> targetDimSizes = dimSizes;
-    std::swap(targetDimSizes[dim1Mapped], targetDimSizes[dim2Mapped]);
-    
-    std::vector<tensorSize_t> targetStrides(numDims);
-    stride = 1;
-    for(int d = numDims - 1; d >= 0; d--) {
-        targetStrides[d] = stride;
-        stride *= targetDimSizes[d];
-    }
-
-    auto transposedValues = make_unique<tensorValues_t>(source.values->getDevice());
-    transposedValues->resize(source.values->getSize());
-
-    tensorSize_t totalSize = source.values->getSize();
-    for(tensorSize_t targetIdx = 0; targetIdx < totalSize; targetIdx++) {
-        // linear target index to multi-dimensional idx
-        tensorSize_t tmp = targetIdx;
-        for(int d = 0; d < numDims; d++) {
-            indices[d] = tmp / targetStrides[d];
-            tmp %= targetStrides[d];
-        }
-        
-        // Swap the transposed dimensions to get source indices
-        std::swap(indices[dim1Mapped], indices[dim2Mapped]);
-        
-        // Convert multi-dimensional source indices to linear index
-        tensorSize_t sourceIdx = 0;
-        for(int d = 0; d < numDims; d++) {
-            sourceIdx += indices[d] * sourceStrides[d];
-        }
-        
-        (*transposedValues)[targetIdx] = (*source.values)[sourceIdx];
-    }
-    
-    target.values = std::move(transposedValues);
-    target.dims.swap(dim1Mapped, dim2Mapped);
-}
-
-/**
- * @brief Transposes 2D tensor. A little sleeker than transposeImpl.
- */
-void Tensor::transposeImpl2D(Tensor& target, const int dim1, const int dim2) const noexcept {
-  assert(values->getSize()==target.values->getSize() && target.dims.nDims()==2 && dims.nDims()==2);
-  if(target.getDevice()==Device::CUDA){
-    __throw_runtime_error("Transposition for CUDA not implemented yet");
-  }
-  else if(dim1==dim2){
-    return;
-  }
-
-  const auto& source = *this; // easier to read
-
-  auto dim1Mapped = mapDim(dim1, source.dims);
-  auto dim2Mapped = mapDim(dim2, source.dims);
-
-  // large dim wraps small dim
-  const auto largeDim = dim1Mapped < dim2Mapped ? dim1Mapped : dim2Mapped;
-  const auto smallDim = dim1Mapped < dim2Mapped ? dim2Mapped : dim1Mapped;
-
-  // largeDimSize >= smallDimSize
-  const auto largeDimOffset = getDimOffset(largeDim, dims);
-  const auto smallDimOffset = getDimOffset(smallDim, dims);
-
-  auto transposedValues = make_unique<tensorValues_t>(source.values->getDevice());
-  transposedValues->resize(source.values->getSize());
-
-  tensorSize_t resIdx = 0;
-  for(tensorSize_t smallDimCount=0; smallDimCount<source.dims.get(smallDim); smallDimCount++){
-    for(tensorSize_t largeDimCount=0; largeDimCount<source.dims.get(largeDim); largeDimCount++){
-      tensorSize_t offset = largeDimCount * largeDimOffset + smallDimCount * smallDimOffset;
-
-      for(tensorSize_t smallDimIdx=0; smallDimIdx<smallDimOffset; smallDimIdx++){
-        (*transposedValues)[resIdx] = (*source.values)[offset];
-        resIdx++;
-        offset++;
-      }
-    }
-  }
-
-  target.values = std::move(transposedValues);
-  target.dims.swap(dim1Mapped, dim2Mapped);
-
-  /* if(source.grads){
-    source.grads->transposeImpl(*target.grads, dim1, dim2); // TODO: do we need this?
-  } */
-}
-
-
-/**
- * @brief Swap dim1 and dim2, modify this tensor.
- * 
- * Out of place operation.
- */
-void Tensor::transposeThis(int dim1, int dim2) noexcept {
-  if(dims.nDims()>2){
-    transposeImpl(*this, dim1, dim2);
-  }
-  else{
-    transposeImpl2D(*this, dim1, dim2);
-  }
-}
-
-/**
- * @brief Out of place transposition of last two axes.
- * 
- */
-void Tensor::transposeThis() noexcept {
-  if(dims.nDims()<2){
-    return;
-  }
-  transposeThis(-1, -2);
-}
-
 /**
- * @brief Like overloaded transpose with requiresGrad==false.
+ * @brief Get a tensor that is linear in memory. Useful for coalesced memory access patterns.
  */
-Tensor Tensor::transpose(int dim1, int dim2) const {
-  return transpose(dim1, dim2, false);
+Tensor Tensor::getContiguous() const {
+  if(dims.inOriginalState())
+    return createShallowCopy();
+  return createContiguousCopy();
 }
 
 /**
- * @brief Like transposeThis, but returns a new tensor.
- * We give requiresGrad as an optional argument to give more control of 
- * what this tensor is intended to do. E.g. in backprop sometimes we do 
- * need to create transposed tensors to multiply with, and with 
- * requiresGrad==false we avoid unnecessary memory allocation overhead.
+ * @brief Quick transpose operation.
  */
-Tensor Tensor::transpose(int dim1, int dim2, const bool requiresGrad) const {
-  Tensor res(dims, values->getDevice(), requiresGrad);
-  
-  if(dims.nDims()>2){
-    transposeImpl(res, dim1, dim2);
-  }
-  else{
-    transposeImpl2D(res, dim1, dim2);
-  }
-
-  return res;
+Tensor Tensor::transpose(int dim1, int dim2) {
+  Tensor result = createShallowCopy();
+  result.dims.swap(dim1, dim2);
+  return result;
 }
 
 /**
@@ -759,11 +848,11 @@ Tensor Tensor::transpose(int dim1, int dim2, const bool requiresGrad) const {
  * 
  * New order aligns axes newly. E.g. (2, 3, 1, 0)
  */
-void Tensor::permute(const std::vector<tensorDim_t>&& newOrder) noexcept {
-  // TODO: highly inefficient -> refactor
-  assert(newOrder.size()<=std::numeric_limits<tensorDim_t>::max());
+void Tensor::permute(const std::vector<tensorDim_t>& newOrder) noexcept {
+  assert(newOrder.size()==dims.nDims());
+
   for(tensorDim_t i=0; i<static_cast<tensorDim_t>(newOrder.size()); i++){
-    transpose(i, newOrder[i]);
+    dims.swap(i, newOrder[i]);
   }
 }
 
@@ -771,17 +860,46 @@ void Tensor::permute(const std::vector<tensorDim_t>&& newOrder) noexcept {
  * @brief Populates the tensor with value.
  */
 void Tensor::reset(const ftype x) noexcept {
-  for(tensorSize_t i=0; i<values->getSize(); i++){
-    (*values)[i] = x;
+  switch(values->getDevice()) {
+    case Device::CPU:
+      std::fill(values->data(), values->data() + values->getSize(), x);
+      break;
+    case Device::CUDA:
+      #ifdef __CUDA
+        cuda_impl::scalarFill(*this, x);
+      #else
+        __throw_runtime_error("Not compiled with CUDA");
+      #endif
+      break;
   }
 }
 
 /**
  * @brief Populates the tensor with values drawn according to initializer.
  */
-void Tensor::reset(const shared_ptr<utility::InitializerBase> init) noexcept {
-  for(tensorSize_t i=0; i<values->getSize(); i++){
-    (*values)[i] = init->drawNumber();
+void Tensor::reset(shared_ptr<utility::InitializerBase> init) noexcept {
+  const auto size = values->getSize();
+  switch(values->getDevice()) {
+    case Device::CPU:
+      for(tensorSize_t i = 0; i < size; i++){
+        values->data()[i] = init->drawNumber();
+      }
+      break;
+    case Device::CUDA:
+      #ifdef __CUDA
+      {
+        auto newValues = static_cast<ftype*>(std::malloc(size * sizeof(ftype)));
+        for(tensorSize_t i = 0; i < size; i++){
+          newValues[i] = init->drawNumber();
+        }
+        cudaErrchk(cudaMemcpy(values->data(), newValues, size * sizeof(ftype), cudaMemcpyHostToDevice));
+        free(newValues);
+        // TODO: better initialize directly on GPU
+      }
+      #else
+        __throw_runtime_error("Not compiled with CUDA");
+      #endif
+      break;
   }
 }
 
@@ -802,7 +920,12 @@ Device Tensor::getDefaultDevice() noexcept {
 }
 
 void Tensor::setDevice(const Device d) noexcept {
+  makeContiguous();
   values->setDevice(d);
+
+  if(grads) {
+    grads->setDevice(d);
+  }
 }
 
 Device Tensor::getDevice() const noexcept {
@@ -819,11 +942,13 @@ Device Tensor::getDevice() const noexcept {
  * @param high Upper idx, non-inclusive bound.
  * @return Tensor The slices tensor.
  */
-Tensor Tensor::getSlice(tensorSize_t low, tensorSize_t high) const {
+Tensor Tensor::getSlice(const tensorSize_t low, const tensorSize_t high) const {
   if(high<=low){
     __throw_invalid_argument("Upper bound most be larger than lower bound.");
   }
 
+  makeContiguous();
+  
   auto resDims = dims.toVector();
   resDims[0] = high-low;
   Tensor res(std::move(resDims), values->getDevice(), false);
@@ -840,58 +965,45 @@ Tensor Tensor::getSlice(tensorSize_t low, tensorSize_t high) const {
  */
 Tensor Tensor::getSlice(span<const tensorDim_t> indices) const {
   assert(indices.size()>0);
-  
+
+  makeContiguous();
   auto resDims = dims.toVector();
   resDims[0] = indices.size();
 
   Tensor res(std::move(resDims), values->getDevice(), false);
-  values->copyValues(*res.values, indices, getDimOffset(0, resDims));
+  values->copyValues(*res.values, indices, res.getDims().getStride(0));
   return res;
 }
 
-/**
- * @brief Prints only sample of up to 2D tensors.
- */
-void printValuesCpu(std::ostream& os, const Tensor& t) {
+ostream& operator<<(ostream& os, const Tensor& t) noexcept {
+  os << "Dims: " << t.getDims();
+  os << "\nDevice: " << DeviceToString(t.values->getDevice());
+  os << "\nrequiresGrad: " << t.requiresGrad << "\n\n";
+  
   auto printVals = [&os](const Tensor& t){
     constexpr auto MAX_IDX = static_cast<tensorDim_t>(10);
 
-    if(t.dims.nDims()==2){
-      for(tensorDim_t i=0; i<min(MAX_IDX, t.dims.get(0)); i++){
-        for(tensorDim_t j=0; j<min(MAX_IDX, t.dims.get(1)); j++){
+    if(t.dims.nDims() == 2){
+      for(tensorDim_t i = 0; i < min(MAX_IDX, t.dims.get(0)); i++){
+        for(tensorDim_t j = 0; j < min(MAX_IDX, t.dims.get(1)); j++){
           os << t.get({i, j}) << " ";
         }
         os << "\n";
       }
     }
     else{
-      for(tensorDim_t i=0; i<min(MAX_IDX, static_cast<tensorDim_t>(t.values->getSize())); i++){
+      for(tensorDim_t i = 0; i < min(MAX_IDX, static_cast<tensorDim_t>(t.values->getSize())); i++){
         os << (*t.values)[i] << " ";
       }
     }
   };
 
+  t.makeContiguous();
   printVals(t);
   if(t.grads){
     os << "\n\nGrads:\n";
     printVals(*t.grads);
   }
-}
-
-ostream& operator<<(ostream& os, const Tensor& t) noexcept {
-  os << "Dims: " << t.getDims();
-  os << "\nDevice: " << DeviceToString(t.values->getDevice());
-  os << "\nrequiresGrad: " << t.requiresGrad << "\n\n";
-
-  switch(t.values->getDevice()){
-    case Device::CPU:
-      printValuesCpu(os, t);
-      break;
-    case Device::CUDA:
-      __throw_invalid_argument("CUDA not supported yet in printing");
-      break;
-  }
-
   return os;
 }
 
@@ -914,47 +1026,21 @@ tensorSize_t Tensor::computeLinearIdx(const std::vector<tensorDim_t>& idx, const
     __throw_invalid_argument("Number of idxs must match number of dimensions.");
   }
   else if(idx.size()==0){
-    return 0; // TODO: this was 1. What is going on here?
+    return 0;
   }
 
-  const auto lastIdx = idx.size()-1;
-  tensorSize_t offsetFactor = dims.get(lastIdx);
-  
-  tensorSize_t res = idx[lastIdx];
-  for(int i=lastIdx-1; i>=0; i--){
-    res += idx[i] * offsetFactor;
-    offsetFactor *= dims.get(i);
+  tensorSize_t res = 0;
+  for(tensorDim_t i=0; i<idx.size(); i++){
+    res += idx[i] * dims.getStride(i);
   }
-
   return res;
 }
 
-/**
- * @brief Gets the total size of a dimension. E.g. if dims=(2, 3, 4),
- * the offset of dim1 is 3*4==12, and that of dim0 is 2*3*4==24.
- */
-tensorSize_t Tensor::getDimOffset(const tensorDim_t dim, const Dimension& dims) {
-  tensorSize_t res = 1; // minimum possible dimsize
-
-  for(size_t idx = dims.nDims()-1; idx>dim; idx--){
-    res *= dims.get(idx);
-  }
-
-  assert(res!=0);
-  return res;
-}
-
-/**
- * @brief Like overload, but accepts negative dims.
- */
-tensorSize_t Tensor::getDimOffset(const int dim, const Dimension& dims) {
-  return getDimOffset(mapDim(dim, dims), dims);
-}
-
 /**
  * @brief No explanation needed.
  */
 ftype Tensor::get(const std::vector<tensorDim_t>& idx) const {
+  makeContiguous();
   return (*values)[computeLinearIdx(idx, dims)]; 
 }
 
@@ -973,7 +1059,6 @@ ftype Tensor::operator[](tensorSize_t idx) const {
   return (*values)[idx];
 }
 
-
 ftype Tensor::get(tensorDim_t idx0, tensorDim_t idx1) const {
   return get({idx0, idx1});
 }
@@ -990,7 +1075,8 @@ ftype Tensor::get(tensorDim_t idx0, tensorDim_t idx1, tensorDim_t idx2, tensorDi
  * @brief No explanation needed.
  */
 void Tensor::set(ftype item, const std::vector<tensorDim_t>& idx) {
-  (*values)[computeLinearIdx(idx, dims)] = item;
+  makeContiguous();
+  values->set(item, computeLinearIdx(idx, dims));
 }
 
 /**
@@ -998,7 +1084,7 @@ void Tensor::set(ftype item, const std::vector<tensorDim_t>& idx) {
  * Can lead to unexpected results in multidimensional tensors.
  */
 void Tensor::set(ftype item, tensorDim_t idx) { 
-  (*values)[idx] = item;
+  values->set(item, idx);
 }
 
 void Tensor::set(ftype item, tensorDim_t idx0, tensorDim_t idx1) { 
diff --git a/src/backend/data_modeling/tensor.h b/src/backend/data_modeling/tensor.h
index 01f4091..b7e6a55 100644
--- a/src/backend/data_modeling/tensor.h
+++ b/src/backend/data_modeling/tensor.h
@@ -40,6 +40,8 @@ class Tensor final : public std::enable_shared_from_this<Tensor>
   friend class cgraph::TopologicalSort;
 
 private:
+  struct shallowCopyToken{}; // we only want createShallowCopy to do shallow copies
+
   /**
    * @brief Here we encapsulate the tensor's values.
    * Enables us to use a shared_ptr, as well as encapsulate all the
@@ -51,13 +53,11 @@ class Tensor final : public std::enable_shared_from_this<Tensor>
   {
   private:
     tensorSize_t size = 0;
-    ftype *values = nullptr;
+    ftype* values = nullptr;
 
     Device device;
     inline static Device defaultDevice = Device::CPU;
 
-    void addOtherCpu(const tensorValues_t& other) noexcept;
-
   public:
     explicit tensorValues_t();
     explicit tensorValues_t(Device d);
@@ -71,94 +71,80 @@ class Tensor final : public std::enable_shared_from_this<Tensor>
 
     void copyFromRaw(const ftype* src, tensorSize_t n);
 
+    ftype* data() noexcept { return values; }
+    const ftype* data() const noexcept { return values; }
+    
     explicit operator bool() const noexcept;
-    ftype& operator[](const tensorSize_t idx);
-    ftype operator[](const tensorSize_t idx) const;
 
+    ftype operator[](tensorSize_t idx) const;
     void set(ftype v, tensorSize_t idx);
     ftype get(tensorSize_t idx);
 
     tensorSize_t getSize() const noexcept;
 
-    // needed for gradient descent
-    tensorValues_t& operator+=(const tensorValues_t& other);
-
-    template<typename T>
-      requires(std::is_integral_v<std::remove_const_t<T>>)
-    void resize(const T size)
-    {
-      this->size = static_cast<tensorSize_t>(size);
-      switch (this->device)
-      {
-      case Device::CPU:
-        values = static_cast<ftype* >(std::malloc(this->size * sizeof(ftype)));
-        break;
-      case Device::CUDA:
-        std::__throw_invalid_argument("Not implemented yet.");
-        break;
-      }
-    }
+    void resize(tensorSize_t size);
 
-    void setDevice(const Device d) noexcept;
+    void setDevice(Device d) noexcept;
     Device getDevice() const noexcept;
 
+
     void copyValues(tensorValues_t& target) const;
     void copyValues(tensorValues_t& target, tensorSize_t low, tensorSize_t high, tensorSize_t targetOffset) const;
-    void copyValues(tensorValues_t& target, std::span<const tensorDim_t> indices, const tensorSize_t sizeOfDim) const;
+    void copyValues(tensorValues_t& target, std::span<const tensorDim_t> indices, tensorSize_t sizeOfDim) const;
 
-    static void setDefaultDevice(const Device d) noexcept;
+    static void setDefaultDevice(Device d) noexcept;
     static Device getDefaultDevice() noexcept;
   };
 
-  Dimension dims;
-  std::unique_ptr<tensorValues_t> values = nullptr; // contained values of tensor
+  mutable Dimension dims;
+  mutable std::shared_ptr<tensorValues_t> values = nullptr; // contained values of tensor
 
   bool requiresGrad = false;
   std::shared_ptr<Tensor> grads = nullptr; // gradients
   std::shared_ptr<cgraph::GraphNode> cgNode = nullptr;
 
-  static Tensor multiplyScalar(const Tensor& scalar, const Tensor& other) noexcept;
-
   static Tensor matMulImpl(const Tensor& left, const Tensor& right);
   static void matMul2DCpu(Tensor& res, const Tensor& left, const Tensor& right,
-                          const tensorSize_t resOffset, const tensorSize_t leftOffset,
-                          const tensorSize_t rightOffset);
+                          tensorSize_t resOffset, tensorSize_t leftOffset,
+                          tensorSize_t rightOffset);
 
-  void transposeImpl2D(Tensor& target, const int dim1, const int dim2) const noexcept;
-  void transposeImpl(Tensor& target, const int dim1, const int dim2) const noexcept;
+  void makeContiguous() const;
 
   // convenience functions that appear in multiple places
   static tensorSize_t computeLinearIdx(const std::vector<tensorDim_t>&& idx, const Dimension& dims);
   static tensorSize_t computeLinearIdx(const std::vector<tensorDim_t>& idx, const Dimension& dims);
 
-  static tensorSize_t getDimOffset(const tensorDim_t dim, const Dimension& dims);
-  static tensorSize_t getDimOffset(const int dim, const Dimension& dims);
-  static tensorDim_t mapDim(const int dim, const Dimension& dims);
+  static tensorDim_t mapDim(int dim, const Dimension& dims);
 
-  friend void printValuesCpu(std::ostream& os, const Tensor& t);
+  Tensor(const Tensor& other, shallowCopyToken);
 
 public:
   template <typename T>
     requires(std::is_same_v<std::remove_cvref_t<T>, Dimension>)
   explicit Tensor(T&& dims, Device d, bool requiresGrad = false)
-      : Tensor{dims.toVector(), tensorValues_t::getDefaultDevice(), requiresGrad}
+      : Tensor{dims.toVector(), d, requiresGrad}
+      // !!!needs dims.toVector() to not trigger the copy ctors!!!
   { }
 
-  explicit Tensor(const std::vector<tensorDim_t>& dims, bool requiresGrad = false) : dims{dims}, values{std::make_unique<tensorValues_t>()}, requiresGrad{requiresGrad}
+  explicit Tensor(const std::vector<tensorDim_t>& dims, bool requiresGrad = false) 
+    : dims{dims}, values{std::make_unique<tensorValues_t>()}, requiresGrad{requiresGrad}
   {
     values->resize(this->dims.getSize());
   }
 
-  explicit Tensor(const std::vector<tensorDim_t>& dims, Device d, bool requiresGrad = false) : dims{dims}, values{std::make_unique<tensorValues_t>(d)}, requiresGrad{requiresGrad}
+  explicit Tensor(const std::vector<tensorDim_t>& dims, Device d, bool requiresGrad = false) 
+    : dims{dims}, values{std::make_unique<tensorValues_t>(d)}, requiresGrad{requiresGrad}
   {
     values->resize(this->dims.getSize());
   }
 
-  explicit Tensor(const std::vector<tensorDim_t>& dims, const std::vector<ftype>& initValues, bool requiresGrad = false) : Tensor{dims, std::move(initValues), Tensor::getDefaultDevice(), requiresGrad}
+  explicit Tensor(const std::vector<tensorDim_t>& dims, const std::vector<ftype>& initValues, bool requiresGrad = false) 
+    : Tensor{dims, std::move(initValues), Tensor::getDefaultDevice(), requiresGrad}
   {
   }
 
-  explicit Tensor(const std::vector<tensorDim_t>& dims, const std::vector<ftype>& initValues, Device d, bool requiresGrad = false) : Tensor{dims, d, requiresGrad}
+  explicit Tensor(const std::vector<tensorDim_t>& dims, const std::vector<ftype>& initValues, Device d, bool requiresGrad = false) 
+    : Tensor{dims, d, requiresGrad}
   {
     for (tensorSize_t i=0; i<initValues.size(); i++){
       values->set(initValues[i], i);
@@ -182,6 +168,8 @@ class Tensor final : public std::enable_shared_from_this<Tensor>
   Tensor& operator=(const Tensor& other) = delete;
 
   Tensor createEmptyCopy() const;
+  Tensor createShallowCopy() const;
+  Tensor createContiguousCopy() const;
   Tensor createDeepCopy() const;
 
   /**
@@ -191,8 +179,10 @@ class Tensor final : public std::enable_shared_from_this<Tensor>
   Tensor(Tensor&& other) noexcept;
   Tensor& operator=(Tensor&& other) noexcept;
 
-  void reset(const ftype x) noexcept;
-  void reset(const std::shared_ptr<utility::InitializerBase> init) noexcept;
+  ftype* getData() const noexcept;
+
+  void reset(ftype x) noexcept;
+  void reset(std::shared_ptr<utility::InitializerBase> init) noexcept;
 
   const Dimension& getDims() const noexcept;
   tensorSize_t getSize() const noexcept;
@@ -208,6 +198,8 @@ class Tensor final : public std::enable_shared_from_this<Tensor>
   Tensor operator*(const Tensor& t) const;
   Tensor elementwiseMul(const Tensor& other) const;
 
+  Tensor& operator+=(const Tensor& other);
+
   // TODO: Tensor operator/(const Tensor& other) const;
 
   // the following operators all broadcast
@@ -228,13 +220,11 @@ class Tensor final : public std::enable_shared_from_this<Tensor>
   }
   bool hasGrads() const noexcept { return grads!=nullptr; }
 
-  void transposeThis() noexcept;
-  void transposeThis(int dim1, int dim2) noexcept;
-
-  Tensor transpose(int dim1, int dim2) const;
-  Tensor transpose(int dim1, int dim2, const bool requiresGrad) const;
+  Tensor transpose(int dim1=-1, int dim2=-2);
+  void permute(const std::vector<tensorDim_t>& newOrder) noexcept;
 
-  void permute(const std::vector<tensorDim_t>&& newOrder) noexcept;
+  bool isContiguous() const noexcept { return dims.inOriginalState(); }
+  Tensor getContiguous() const;
 
   friend std::ostream& operator<<(std::ostream& os, const Tensor& t) noexcept;
 
diff --git a/src/backend/data_modeling/tensor_functions.cpp b/src/backend/data_modeling/tensor_functions.cpp
index 5f06d2d..2744f92 100644
--- a/src/backend/data_modeling/tensor_functions.cpp
+++ b/src/backend/data_modeling/tensor_functions.cpp
@@ -11,6 +11,11 @@
 
 #include "tensor_functions.h"
 
+#ifdef __CUDA
+#include "data_modeling/cuda/tensor_ops.cuh"
+#include "computational_graph/tensor_ops/cuda/tensor_ops_nodes.cuh"
+#endif
+
 using namespace std;
 
 Tensor TensorFunctions::Zeros(vector<tensorDim_t> dims, Device d, const bool requiresGrad) {
@@ -33,8 +38,8 @@ Tensor TensorFunctions::Ones(vector<tensorDim_t> dims, const bool requiresGrad)
   return Ones(std::move(dims), Tensor::getDefaultDevice(), requiresGrad);
 }
 
-Tensor TensorFunctions::Gaussian(vector<tensorDim_t> dims, const Device d, 
-                                 const ftype stddev, const bool requiresGrad) {
+Tensor TensorFunctions::Gaussian(vector<tensorDim_t> dims, const ftype stddev,  
+                                 const Device d, const bool requiresGrad) {
   auto res = Tensor(std::move(dims), d, requiresGrad);
   res.reset(std::make_shared<utility::GaussianInitializer>(stddev));
   return res;
@@ -42,7 +47,7 @@ Tensor TensorFunctions::Gaussian(vector<tensorDim_t> dims, const Device d,
     
 Tensor TensorFunctions::Gaussian(vector<tensorDim_t> dims, const ftype stddev, 
                                  const bool requiresGrad) {
-  return Gaussian(std::move(dims), Tensor::getDefaultDevice(), stddev, requiresGrad);
+  return Gaussian(std::move(dims), stddev, Tensor::getDefaultDevice(), requiresGrad);
 }
 
 // Tensor manipulation
@@ -90,24 +95,46 @@ shared_ptr<Tensor> TensorFunctions::makeSharedTensor(const vector<tensorDim_t>&
   * Input dim must be smaller then t.dims.nDims()-1
   */
 Tensor TensorFunctions::SumOverDims(const Tensor& t, tensorDim_t dim) {
-  if(dim>=t.getDims().nDims()-1){
+  if(dim > t.getDims().nDims() - 1){
     __throw_invalid_argument("Dim parameter must be smaller than number of dims, but was " + dim);
   }
 
   auto resDims = t.getDims().collapseDimension(dim);
   Tensor res = Zeros(resDims.toVector(), t.getDevice(), t.getRequiresGrad()); // inefficiency toVector
 
-  tensorSize_t stride = 1;
-  for(tensorDim_t i=dim+1; i<t.getDims().nDims(); i++){
-    stride *= t.getDims()[i];
-  }
-  
-  tensorSize_t targetOffset = 0;
-  for(tensorDim_t loop=0; loop<t.getDims()[dim]; loop++){
-    for(tensorSize_t i=0; i<stride; i++){
-      res.set(res.get(i) + t.get(targetOffset), i);
-      targetOffset++;
+  switch(t.getDevice()) {
+    case Device::CPU:
+    {
+      tensorSize_t stride = 1;
+      for(tensorDim_t i = dim + 1; i < t.getDims().nDims(); i++){
+        stride *= t.getDims()[i];
+      }
+
+      // size of dimensions before dim
+      tensorSize_t outerSize = 1;
+      for(tensorDim_t i = 0; i < dim; i++) {
+        outerSize *= t.getDims()[i];
+      }
+      
+      for(tensorSize_t outer = 0; outer < outerSize; outer++){
+        for(tensorDim_t k = 0; k < t.getDims()[dim]; k++){
+          for(tensorSize_t i = 0; i < stride; i++){
+            tensorSize_t srcIdx = outer * t.getDims()[dim] * stride + k * stride + i;
+            tensorSize_t dstIdx = outer * stride + i;
+            
+            res.set(res.get(dstIdx) + t.get(srcIdx), dstIdx);
+          }
+        }
+      }
+      break;
     }
+    case Device::CUDA:
+      #ifdef __CUDA
+        cuda_impl::sumOverDims(res, t, dim);
+      #else
+        __throw_invalid_argument("Not compiled with CUDA");
+      #endif
+      break;
   }
   
   return res;
diff --git a/src/backend/data_modeling/tensor_functions.h b/src/backend/data_modeling/tensor_functions.h
index 2d93811..c5fdf0e 100644
--- a/src/backend/data_modeling/tensor_functions.h
+++ b/src/backend/data_modeling/tensor_functions.h
@@ -26,14 +26,14 @@
  */
 namespace TensorFunctions { // class name acts as namespace for us
   // Tensor creation
-  Tensor Zeros(std::vector<tensorDim_t> dims, Device d, const bool requiresGrad=false);
-  Tensor Zeros(std::vector<tensorDim_t> dims, const bool requiresGrad=false);
+  Tensor Zeros(std::vector<tensorDim_t> dims, Device d, bool requiresGrad=false);
+  Tensor Zeros(std::vector<tensorDim_t> dims, bool requiresGrad=false);
 
-  Tensor Ones(std::vector<tensorDim_t> dims, Device d, const bool requiresGrad=false);
-  Tensor Ones(std::vector<tensorDim_t> dims, const bool requiresGrad=false);
+  Tensor Ones(std::vector<tensorDim_t> dims, Device d, bool requiresGrad=false);
+  Tensor Ones(std::vector<tensorDim_t> dims, bool requiresGrad=false);
 
-  Tensor Gaussian(std::vector<tensorDim_t> dims, Device d, ftype stddev, const bool requiresGrad=false);
-  Tensor Gaussian(std::vector<tensorDim_t> dims, ftype stddev=1, const bool requiresGrad=false);
+  Tensor Gaussian(std::vector<tensorDim_t> dims, ftype stddev, Device d, bool requiresGrad=false);
+  Tensor Gaussian(std::vector<tensorDim_t> dims, ftype stddev=1, bool requiresGrad=false);
 
   std::shared_ptr<Tensor> makeSharedTensor(const std::vector<tensorDim_t>& dims, bool requiresGrad=false);
 
diff --git a/src/backend/module/activation_functions/cuda/activations.cu b/src/backend/module/activation_functions/cuda/activations.cu
new file mode 100644
index 0000000..889350c
--- /dev/null
+++ b/src/backend/module/activation_functions/cuda/activations.cu
@@ -0,0 +1,464 @@
+/**
+ * @file activations.cu
+ * @author Robert Baumgartner (r.baumgartner-1@tudelft.nl)
+ * @brief 
+ * @version 0.1
+ * @date 2026-03-31
+ * 
+ * @copyright Copyright (c) 2026
+ * 
+ */
+
+#ifndef __CUDA
+static_assert(false, "File should not be compiled without CUDA enabled");
+#endif // __CUDA
+
+#include "activations.cuh"
+
+#include "shared/cuda/common_kernels.cuh"
+#include "shared/cuda/common_softmax.cuh"
+
+#include "utility/utils.h"
+#include "utility/cuda/cuda_common.cuh"
+
+#include <stdexcept>
+
+using namespace std;
+
+namespace {
+  using namespace cuda_impl;
+
+  /**
+   * @brief Kernel for forward ReLU function.
+   */
+  __global__ void reluKernel(ftype* const res, const ftype* const input, const tensorSize_t size) {
+    const int gid = blockIdx.x * blockDim.x + threadIdx.x;
+    if(gid >= size)
+      return;
+
+    res[gid] = cudaMax<ftype>(input[gid], 0);
+  }
+
+  /**
+   * @brief Kernel for forward Leaky-ReLU function.
+   */
+  __global__ void leakyReluKernel(ftype* const res, const ftype* const input, const ftype eps, const tensorSize_t size) {
+    const int gid = blockIdx.x * blockDim.x + threadIdx.x;
+    if(gid >= size)
+      return;
+
+    res[gid] = cudaMax<ftype>(input[gid], eps * input[gid]); // eps < 1
+  }
+
+  /**
+   * @brief Kernel for forward Sigmoid function.
+   */
+  __global__ void sigmoidKernel(ftype* const res, const ftype* const input, const tensorSize_t size) {
+    const int gid = blockIdx.x * blockDim.x + threadIdx.x;
+    if(gid >= size)
+      return;
+
+    res[gid] = cudaSigmoid<ftype>(input[gid]);
+  }
+
+  /**
+   * @brief Reduction kernel that computes the sum over an array within the size of 2 * warpsize at maximum.
+   */
+  template<int maxoffset>
+  __forceinline__ __device__ void softmaxWarpSumReduce(volatile ftype* const input, const tensorSize_t stride, const int offset) {
+    // TODO: warp shuffle for newer architectures
+    if(maxoffset == 32) {
+      if(offset + 32 < stride) input[offset] += input[offset + 32];
+    }
+    if(maxoffset >= 16) {
+      if(offset + 16 < stride) input[offset] += input[offset + 16];
+    }
+    if(maxoffset >= 8) {
+      if(offset + 8 < stride) input[offset] += input[offset + 8];
+    }
+    if(maxoffset >= 4) {
+      if(offset + 4 < stride) input[offset] += input[offset + 4];
+    }
+    if(maxoffset >= 2) {
+      if(offset + 2 < stride) input[offset] += input[offset + 2];
+    }
+    if(maxoffset >= 1) {
+      if(offset + 1 < stride) input[offset] += input[offset + 1];
+    }
+  }
+
+  /**
+   * @brief Numerically stable version of softmax kernel. Just as in findMaxKernelOneWarp we assume that stride <= 2 * warpsize.
+   * Numerical stability comes from computing the maximum values per row, see findMaxKernelOneWarp and argument maxValues.
+   */
+  template<typename T, int maxoffset>
+  __global__ void stableSoftmaxKernelOneWarp(ftype* const res, const ftype* const input, const ftype* const maxValues,
+                                             const tensorSize_t stride, const tensorSize_t size) {
+    assert(blockDim.x % 32 == 0);
+
+    int gid = blockIdx.x * blockDim.x + threadIdx.x;
+    if(gid >= size)
+      return;
+
+    const auto strideOffset = gid / stride;
+    const auto maxValue = maxValues[strideOffset];
+    ftype expVal = stableExp<ftype>(input[gid], maxValue);
+
+    int tid = threadIdx.x;
+    extern __shared__ ftype smem[]; // can lead to bank conflicts iff std::is_same_v<T, double>
+    smem[tid] = expVal;
+    __syncthreads();
+
+    volatile ftype* const start = smem + (tid / stride) * stride;
+    const int offset = gid % stride;
+    softmaxWarpSumReduce<maxoffset>(start, stride, offset);
+
+    res[gid] = expVal / start[0];
+  }
+
+  /**
+   * @brief Just like stableSoftmaxKernelOneWarp, but this one works across a whole block, not just a warp.
+   * 
+   * In this initial version we assume one kernel per stride, to make matters simple to understand.
+   */
+  template<typename T>
+  __global__ void stableSoftmaxKernelOneBlock(ftype* const res, const ftype* const input, const ftype* const maxValues, const tensorSize_t stride) {
+    assert_debug(blockDim.x / stride == 0, "Kernel built for one stride per block, blockDim.x is < stride"); 
+
+    const int tid = threadIdx.x;
+    const int gid = blockIdx.x * stride + tid;
+
+    extern __shared__ ftype smem[]; // can lead to bank conflicts iff std::is_same_v<T, double>
+    const auto maxValue = maxValues[blockIdx.x];
+    ftype expVal = stableExp<ftype>(input[gid], maxValue);
+    smem[tid] = expVal;
+
+    const tensorSize_t maxIdx = tid + blockDim.x;
+    const bool doPadding = maxIdx >= stride;
+
+    ftype expValOffset = 0;
+    if(!doPadding) { // some threads will be idle here
+      expValOffset = stableExp<ftype>(input[gid + blockDim.x], maxValue);
+    }
+
+    smem[maxIdx] = expValOffset;
+    __syncthreads();
+
+    for(tensorSize_t offset = blockDim.x; offset > 32; offset >>= 1) {
+      if(tid < offset) {
+        smem[tid] += smem[tid + offset];
+      }
+      __syncthreads();
+    }
+
+    // TODO: warp shuffle for newer architectures
+    volatile ftype* const start = smem;
+    if(tid < 32) {
+      start[tid] += start[tid + 32];
+      start[tid] += start[tid + 16];
+      start[tid] += start[tid + 8];
+      start[tid] += start[tid + 4];
+      start[tid] += start[tid + 2];
+      start[tid] += start[tid + 1];
+    }
+    __syncthreads(); // needed because threads > 32 will also use start[0]
+
+    res[gid] = expVal / start[0];
+    if(!doPadding) {
+      res[gid + blockDim.x] = expValOffset / start[0];
+    }
+  }
+
+  /**
+   * @brief Does the first part of stableSoftmaxKernelOneBlock, namely the sums. Because here again we have a partial sum
+   * and assume the stride did not fit into the block entirely, we do a partial sum only. Additionally, write the max-adjusted
+   * exp values back to res to prepare for the division.
+   */
+  template<typename T>
+  __global__ void stableSoftmaxLargePass1(ftype* const res, ftype* const partialSums, const ftype* const input, const ftype* const maxValues, 
+                                          const tensorSize_t stride, const int blocksPerStride) {
+    const int tid = threadIdx.x;
+    const int strideIdx = blockIdx.x / blocksPerStride;
+    const int blockWithinStride = blockIdx.x % blocksPerStride;
+
+    // same logic as in findMaxKernelLargePass1
+    const int inputBase = strideIdx * stride + blockWithinStride * 2 * blockDim.x;
+    const ftype maxValue = maxValues[strideIdx];
+    
+    extern __shared__ ftype smem[];
+    const tensorSize_t localIdx0 = inputBase + tid;
+    const tensorSize_t localIdx1 = inputBase + tid + blockDim.x;
+    
+    // same logic as in findMaxKernelLargePass1
+    ftype expVal0 = (localIdx0 < (strideIdx + 1) * stride) ? stableExp<T>(input[localIdx0], maxValue) : 0.0f;
+    ftype expVal1 = (localIdx1 < (strideIdx + 1) * stride) ? stableExp<T>(input[localIdx1], maxValue) : 0.0f;
+    
+    smem[tid] = expVal0;
+    smem[tid + blockDim.x] = expVal1;
+    __syncthreads();
+    
+    // write values to output -> will be nominator in division
+    if(localIdx0 < (strideIdx + 1) * stride) {
+      res[localIdx0] = expVal0;
+    }
+    if(localIdx1 < (strideIdx + 1) * stride) {
+      res[localIdx1] = expVal1;
+    }
+    
+    // reduce sum
+    for(tensorSize_t offset = blockDim.x; offset > 32; offset >>= 1) {
+      if(tid < offset) {
+        smem[tid] += smem[tid + offset];
+      }
+      __syncthreads();
+    }
+
+    volatile ftype* start = smem;
+    if(tid < 32) {
+      start[tid] += start[tid + 32];
+      start[tid] += start[tid + 16];
+      start[tid] += start[tid + 8];
+      start[tid] += start[tid + 4];
+      start[tid] += start[tid + 2];
+      start[tid] += start[tid + 1];
+    }
+
+    if(tid == 0) {
+      partialSums[blockIdx.x] = start[0];
+    }
+  }
+
+  /**
+   * @brief Self explanatory after stableSoftmaxLargePass1. Continues the sum reduce, does not need to write further to res, since 
+   * pass 1 already did that for us.
+   */
+  __global__ void stableSoftmaxLargePass2(ftype* const sums, const ftype* const partialSums, const tensorSize_t blocksPerStride) {
+    assert_debug(blockDim.x / blocksPerStride == 0, "Kernel built for one stride per block, blockDim.x is < stride"); 
+
+    const int tid = threadIdx.x;
+    const int gid = blockIdx.x * blocksPerStride + tid;
+
+    extern __shared__ ftype smem[]; // can lead to bank conflicts iff std::is_same_v<T, double>
+
+    const tensorSize_t maxIdx = tid + blockDim.x;
+    const bool doPadding = maxIdx >= blocksPerStride;
+    if(doPadding) {
+      smem[tid] = partialSums[gid];
+    }
+    else {
+      smem[tid] = partialSums[gid] + partialSums[gid + blockDim.x];
+    }
+    __syncthreads();
+
+    for(tensorSize_t offset = blockDim.x >> 1; offset > 32; offset >>= 1) {
+      if(tid < offset) {
+        smem[tid] += smem[tid + offset];
+      }
+      __syncthreads();
+    }
+
+    // TODO: warp shuffle for newer architectures
+    volatile ftype* const start = smem;
+    if(tid < 32) {
+      if(tid + 32 < blockDim.x) start[tid] += start[tid + 32];
+      if(tid + 16 < blockDim.x) start[tid] += start[tid + 16];
+      if(tid + 8  < blockDim.x) start[tid] += start[tid + 8];
+      if(tid + 4  < blockDim.x) start[tid] += start[tid + 4];
+      if(tid + 2  < blockDim.x) start[tid] += start[tid + 2];
+      if(tid + 1  < blockDim.x) start[tid] += start[tid + 1];
+    }
+
+    if(tid == 0) { // one block per stride
+      sums[blockIdx.x] = start[0];
+    }
+  }
+
+  /**
+   * @brief Simple kernel doing the division of softmax in the case of a large stride.
+   */
+  __global__ void divideKernel(ftype* const res, const ftype* const sums, const tensorSize_t stride, const tensorSize_t size) {
+    const int gid = blockIdx.x * blockDim.x + threadIdx.x;
+    if(gid >= size) {
+      return;
+    }
+
+    res[gid] /= sums[gid / stride];
+  }
+}
+
+namespace cuda_impl {
+  void relu(Tensor& res, const Tensor& in) {
+    constexpr int threadsPerBlock = 256;
+    const int blocks = (in.getSize() + threadsPerBlock - 1) / threadsPerBlock;
+
+    reluKernel<<<blocks, threadsPerBlock>>>(res.getData(), in.getData(), in.getSize());
+    cudaErrchk(cudaDeviceSynchronize());
+  }
+
+  void leakyRelu(Tensor& res, const Tensor& in, ftype eps) {
+    constexpr int threadsPerBlock = 256;
+    const int blocks = (in.getSize() + threadsPerBlock - 1) / threadsPerBlock;
+
+    leakyReluKernel<<<blocks, threadsPerBlock>>>(res.getData(), in.getData(), eps, in.getSize());
+    cudaErrchk(cudaDeviceSynchronize());
+  }
+
+  void sigmoid(Tensor& res, const Tensor& in) {
+    constexpr int threadsPerBlock = 256;
+    const int blocks = (in.getSize() + threadsPerBlock - 1) / threadsPerBlock;
+
+    sigmoidKernel<<<blocks, threadsPerBlock>>>(res.getData(), in.getData(), in.getSize());
+    cudaErrchk(cudaDeviceSynchronize());
+  }
+
+  /**
+   * @brief Does the softmax computation. Warning: Current implementation can only handle a stride of 
+   * at max 512 * 512 = 262144 floating point numbers. If number exceeds this an exception is throws.
+   * 
+   * For simplicity and for reasons of CUDA efficiency this function is split into 3 segments. 
+   * 1. stride <= 64 -> we use warp level.
+   * 2. stride > 64 && stride < 512 -> we can fit one stride into one block.
+   * 3. stride > 512 -> we have to use two-stage kernel cascading for parallel reduction.
+   */
+  void softmax(Tensor& res, const Tensor& in) {
+    const tensorSize_t stride = static_cast<tensorSize_t>(in.getDims().get(-1));
+    if(stride == 1) {
+      __throw_invalid_argument("Attemped softmax of one element");
+    }
+
+    const int nStrides = in.getSize() / stride;
+
+    // TODO: use some static struct here to prevent this guy from keeping on re-allocating memory
+    ftype* maxValues;
+    const tensorSize_t nMaxValues = nStrides;
+    cudaErrchk(cudaMalloc(&maxValues, nMaxValues * sizeof(ftype)));
+
+    static const auto warpSizeT2 = 2 * DeviceProperties::getWarpSize(); // TODO: can this be a problem in a multi-GPU setting?
+    if(stride <= warpSizeT2) {
+      assert(DeviceProperties::getWarpSize() == 32);
+
+      const int threadsPerBlock = 256;
+      const int blocks = (in.getSize() + threadsPerBlock - 1) / threadsPerBlock;
+
+      if(stride == 2) {
+        findMaxKernelOneWarp<1> <<<blocks, threadsPerBlock, threadsPerBlock * sizeof(ftype)>>>(maxValues, in.getData(), stride, in.getSize());
+      }
+      else if(stride <= 4) {
+        findMaxKernelOneWarp<2> <<<blocks, threadsPerBlock, threadsPerBlock * sizeof(ftype)>>>(maxValues, in.getData(), stride, in.getSize());
+      }
+      else if(stride <= 8) {
+        findMaxKernelOneWarp<4> <<<blocks, threadsPerBlock, threadsPerBlock * sizeof(ftype)>>>(maxValues, in.getData(), stride, in.getSize());
+      }
+      else if(stride <= 16) {
+        findMaxKernelOneWarp<8> <<<blocks, threadsPerBlock, threadsPerBlock * sizeof(ftype)>>>(maxValues, in.getData(), stride, in.getSize());
+      }
+      else if(stride <= 32) {
+        findMaxKernelOneWarp<16> <<<blocks, threadsPerBlock, threadsPerBlock * sizeof(ftype)>>>(maxValues, in.getData(), stride, in.getSize());
+      }
+      else if(stride <= 64) {
+        findMaxKernelOneWarp<32> <<<blocks, threadsPerBlock, threadsPerBlock * sizeof(ftype)>>>(maxValues, in.getData(), stride, in.getSize());
+      }
+      cudaErrchk(cudaDeviceSynchronize());
+
+      if(stride == 2) {
+        stableSoftmaxKernelOneWarp<ftype, 1> <<<blocks, threadsPerBlock, threadsPerBlock * sizeof(ftype)>>>
+                                                (res.getData(), in.getData(), maxValues, stride, in.getSize());
+      }
+      else if(stride <= 4) {
+        stableSoftmaxKernelOneWarp<ftype, 2> <<<blocks, threadsPerBlock, threadsPerBlock * sizeof(ftype)>>>
+                                                (res.getData(), in.getData(), maxValues, stride, in.getSize());
+      }
+      else if(stride <= 8) {
+        stableSoftmaxKernelOneWarp<ftype, 4> <<<blocks, threadsPerBlock, threadsPerBlock * sizeof(ftype)>>>
+                                                (res.getData(), in.getData(), maxValues, stride, in.getSize());
+      }
+      else if(stride <= 16) {
+        stableSoftmaxKernelOneWarp<ftype, 8> <<<blocks, threadsPerBlock, threadsPerBlock * sizeof(ftype)>>>
+                                                (res.getData(), in.getData(), maxValues, stride, in.getSize());
+      }
+      else if(stride <= 32) {
+        stableSoftmaxKernelOneWarp<ftype, 16> <<<blocks, threadsPerBlock, threadsPerBlock * sizeof(ftype)>>>
+                                                (res.getData(), in.getData(), maxValues, stride, in.getSize());
+      }
+      else if(stride <= 64) {
+        stableSoftmaxKernelOneWarp<ftype, 32> <<<blocks, threadsPerBlock, threadsPerBlock * sizeof(ftype)>>>
+                                                (res.getData(), in.getData(), maxValues, stride, in.getSize());
+      }
+      cudaErrchk(cudaDeviceSynchronize());
+    }
+    else if (stride <= 512) {
+      constexpr int maxThreadsPerBlock = 256;
+      assert_debug(stride <= 2 * maxThreadsPerBlock, "If you adapt maxThreadsPerBlock you also have to adapt the limit in the if-clause above");
+
+      constexpr int stridesPerBlock = 1; // if multiple strides per block allowed (adapt kernels!): max(maxThreadsPerBlock / stride, 1);
+
+      // threads per block needs to be power of 2 for reduction to resolve cleanly
+      int threadsPerBlock = 1;
+      while(threadsPerBlock < stride) threadsPerBlock <<= 1;
+      threadsPerBlock /= 2;
+
+      const int blocks = (nStrides + stridesPerBlock - 1) / stridesPerBlock; // gerneralized version iff multiple strides per block allowed
+
+      findMaxKernelOneBlock<<<blocks, threadsPerBlock, 2 * threadsPerBlock * sizeof(ftype)>>>(maxValues, in.getData(), stride);
+      cudaErrchk(cudaDeviceSynchronize());
+
+      stableSoftmaxKernelOneBlock<ftype><<<blocks, threadsPerBlock, 2 * threadsPerBlock * sizeof(ftype)>>>(res.getData(), in.getData(), maxValues, 
+                                                                                                           stride);
+      cudaErrchk(cudaDeviceSynchronize());
+    }
+    else {
+      // stride does not fit into one block. We employ a 2 pass system, where pass one does a partial 
+      // reduction, and pass two does a reduction over the partial reductions.
+    
+      // each block handles up to 512 elements (2 * 256 threads)
+      constexpr int maxThreadsPerBlock = 256;
+      constexpr int elemsPerBlock = 2 * maxThreadsPerBlock; // constant folding
+      const int blocksPerStride = (stride + elemsPerBlock - 1) / elemsPerBlock;
+      assert_debug(blocksPerStride <= 512, "Stride too large for two-pass reduction");
+      
+      const int totalBlocks = nStrides * blocksPerStride;
+
+      // intermediate max values: one per block per stride
+      ftype* partialMaxValues;
+      const tensorSize_t nPartialMax = totalBlocks;
+      cudaErrchk(cudaMalloc(&maxValues, nStrides * sizeof(ftype)));
+      cudaErrchk(cudaMalloc(&partialMaxValues, nPartialMax * sizeof(ftype)));
+
+      // pass 1: reduce each chunk of 512 elements to one partial max
+      // launch blocksPerStride blocks per stride
+      findMaxKernelLargePass1<<<totalBlocks, maxThreadsPerBlock, 2 * maxThreadsPerBlock * sizeof(ftype)>>>(
+                                partialMaxValues, in.getData(), stride, blocksPerStride);
+      cudaErrchk(cudaDeviceSynchronize());
+
+      // pass 2: reduce partial maxes to one max per stride
+      int threadsPass2 = 1;
+      while(threadsPass2 < blocksPerStride) threadsPass2 <<= 1; // threadsPass2 needs to be power of 2
+      threadsPass2 = max(1, threadsPass2 / 2);
+      
+      findMaxKernelLargePass2<<<nStrides, threadsPass2, threadsPass2 * sizeof(ftype)>>>(maxValues, partialMaxValues, blocksPerStride);
+      cudaErrchk(cudaDeviceSynchronize());
+
+      // same two-pass structure for the sum
+      ftype* partialSums;
+      cudaErrchk(cudaMalloc(&partialSums, nPartialMax * sizeof(ftype)));
+
+      stableSoftmaxLargePass1<ftype><<<totalBlocks, maxThreadsPerBlock, 2 * maxThreadsPerBlock * sizeof(ftype)>>>(
+                        res.getData(), partialSums, in.getData(), maxValues, stride, blocksPerStride);
+      cudaErrchk(cudaDeviceSynchronize());
+
+      // pass 2: reduce partial sums
+      stableSoftmaxLargePass2<<<nStrides, threadsPass2, threadsPass2 * sizeof(ftype)>>>(partialSums, partialSums, blocksPerStride);
+      cudaErrchk(cudaDeviceSynchronize());
+
+      // final division pass: divide each exp value by its stride's sum
+      const int nBlocksDivision = (in.getSize() + maxThreadsPerBlock - 1) / maxThreadsPerBlock;
+      divideKernel<<<nBlocksDivision, maxThreadsPerBlock>>>(res.getData(), partialSums, stride, in.getSize());
+      cudaErrchk(cudaDeviceSynchronize());
+
+      cudaErrchk(cudaFree(partialMaxValues));
+      cudaErrchk(cudaFree(partialSums));
+    }
+
+    cudaErrchk(cudaFree(maxValues));
+  }
+}
\ No newline at end of file
diff --git a/src/backend/module/activation_functions/cuda/activations.cuh b/src/backend/module/activation_functions/cuda/activations.cuh
new file mode 100644
index 0000000..6bfd89e
--- /dev/null
+++ b/src/backend/module/activation_functions/cuda/activations.cuh
@@ -0,0 +1,27 @@
+/**
+ * @file activations.cuh
+ * @author Robert Baumgartner (r.baumgartner-1@tudelft.nl)
+ * @brief 
+ * @version 0.1
+ * @date 2026-03-31
+ * 
+ * @copyright Copyright (c) 2026
+ * 
+ */
+
+#pragma once
+
+#ifndef __CUDA
+static_assert(false, "File should not be included without CUDA enabled");
+#endif // __CUDA
+
+#include "utility/global_params.h"
+#include "data_modeling/tensor.h"
+
+namespace cuda_impl {
+  void relu(Tensor& res, const Tensor& in);
+  void leakyRelu(Tensor& res, const Tensor& in, ftype eps);
+
+  void sigmoid(Tensor& res, const Tensor& in);
+  void softmax(Tensor& res, const Tensor& in);
+}
diff --git a/src/backend/module/activation_functions/leaky_relu.cpp b/src/backend/module/activation_functions/leaky_relu.cpp
index 687e2bf..3f46e5a 100644
--- a/src/backend/module/activation_functions/leaky_relu.cpp
+++ b/src/backend/module/activation_functions/leaky_relu.cpp
@@ -12,17 +12,31 @@
 #include "leaky_relu.h"
 #include "computational_graph/activation_functions/leaky_relu_node.h"
 
+#ifdef __CUDA
+#include "module/activation_functions/cuda/activations.cuh"
+#else
+#include <stdexcept>
+#endif
+
 using namespace std;
 using namespace module;
 
 Tensor LeakyReLu::operator()(const Tensor& t) const {
   auto res = t.createDeepCopy();
 
-  for(tensorSize_t i=0; i<t.getSize(); i++){
-    constexpr ftype zero = 0;
-    if(t[i] < zero){
-      res.set(eps, i);
-    }
+  switch(t.getDevice()) {
+    case Device::CPU:
+      for(tensorSize_t i=0; i<t.getSize(); i++){
+        res.set(std::max(t[i], t[i]*eps), i);
+      }
+      break;
+    case Device::CUDA:
+    #ifdef __CUDA
+      cuda_impl::leakyRelu(res, t, eps);
+    #else
+      __throw_invalid_argument("Attempted to give CUDA tensor");
+    #endif
+      break;
   }
 
   return res;
diff --git a/src/backend/module/activation_functions/leaky_relu.h b/src/backend/module/activation_functions/leaky_relu.h
index 94f0c3d..6d6e7c3 100644
--- a/src/backend/module/activation_functions/leaky_relu.h
+++ b/src/backend/module/activation_functions/leaky_relu.h
@@ -19,7 +19,7 @@ namespace module {
       const ftype eps;
 
     public:
-      LeakyReLu(ftype eps=0.01) : eps{eps}
+      LeakyReLu(ftype eps=EPS_LEAKY_RELU) : eps{eps}
       { }
 
       Tensor operator()(const Tensor& t) const override;
diff --git a/src/backend/module/activation_functions/relu.cpp b/src/backend/module/activation_functions/relu.cpp
index 7935304..b93a0a1 100644
--- a/src/backend/module/activation_functions/relu.cpp
+++ b/src/backend/module/activation_functions/relu.cpp
@@ -12,19 +12,38 @@
 #include "relu.h"
 #include "computational_graph/activation_functions/relu_node.h"
 
+#ifdef __CUDA
+#include "module/activation_functions/cuda/activations.cuh"
+#else
+#include <stdexcept>
+#endif
+
 using namespace std;
 using namespace module;
 
 Tensor ReLu::operator()(const Tensor& t) const {
   auto res = t.createDeepCopy();
 
-  for(tensorSize_t i=0; i<t.getSize(); i++){
-    constexpr ftype zero = 0;
-    if(t[i] < zero){
-      res.set(0, i);
-    }
+  switch(t.getDevice()) {
+    case Device::CPU:
+      for(tensorSize_t i=0; i<t.getSize(); i++){
+        constexpr ftype zero = 0;
+        if(t[i] < zero){
+          res.set(zero, i);
+        }
+      }
+      break;
+    case Device::CUDA:
+    #ifdef __CUDA
+      cuda_impl::relu(res, t);
+    #else 
+      __throw_invalid_argument("Attempted to give CUDA tensor");
+    #endif
+      break;
   }
 
+
+
   return res;
 }
 
diff --git a/src/backend/module/activation_functions/sigmoid.cpp b/src/backend/module/activation_functions/sigmoid.cpp
index 765b44a..11be184 100644
--- a/src/backend/module/activation_functions/sigmoid.cpp
+++ b/src/backend/module/activation_functions/sigmoid.cpp
@@ -10,11 +10,16 @@
  */
 
 #include "sigmoid.h"
-
 #include "computational_graph/activation_functions/sigmoid_node.h"
 
 #include <cmath>
 
+#ifdef __CUDA
+#include "module/activation_functions/cuda/activations.cuh"
+#else
+#include <stdexcept>
+#endif
+
 using namespace std;
 using namespace module;
 
@@ -24,17 +29,29 @@ using namespace module;
 Tensor Sigmoid::operator()(const Tensor& t) const {
   auto res = t.createEmptyCopy();
 
-  constexpr ftype one = 1.0;
-  auto compute = [](ftype x){
-    if(x>=0){
-      return one / (one + exp(-x));
-    }
-    auto e = exp(x);
-    return e / (one + e);
-  };
+  switch(t.getDevice()) {
+    case Device::CPU: {
+      constexpr ftype one = 1.0;
+      auto compute = [](ftype x){
+        if(x >= 0){
+          return one / (one + exp(-x));
+        }
+        auto e = exp(x);
+        return e / (one + e);
+      };
 
-  for(tensorSize_t i=0; i<t.getSize(); i++){
-    res.set(compute(t[i]), i);
+      for(tensorSize_t i=0; i<t.getSize(); i++){
+        res.set(compute(t[i]), i);
+      }
+      break;
+    }
+    case Device::CUDA:
+    #ifdef __CUDA
+      cuda_impl::sigmoid(res, t);
+    #else
+      __throw_invalid_argument("Attempted to give CUDA tensor");
+    #endif
+      break;
   }
 
   return res;
diff --git a/src/backend/module/activation_functions/softmax.cpp b/src/backend/module/activation_functions/softmax.cpp
index a001a69..9833cc2 100644
--- a/src/backend/module/activation_functions/softmax.cpp
+++ b/src/backend/module/activation_functions/softmax.cpp
@@ -10,11 +10,16 @@
  */
 
 #include "softmax.h"
-
 #include "computational_graph/activation_functions/softmax_node.h"
 
 #include <cmath>
 
+#ifdef __CUDA
+#include "module/activation_functions/cuda/activations.cuh"
+#else
+#include <stdexcept>
+#endif
+
 using namespace std;
 using namespace module;
 
@@ -28,42 +33,53 @@ Tensor Softmax::operator()(const Tensor& t) const {
     __throw_invalid_argument("Softmax expects input shape of minimum two dimensions");
   }
 
-  const auto nRows = t.getDims()[-2];
-  const auto nCols = t.getDims()[-1];
-
-  // pre-compute exponents
-  Tensor tmp(t.getDims(), t.getDevice(), false);
-  for(tensorDim_t i=0; i<nRows; i++){
-    // for numerical stability, avoid large values
-    // by centering around maxValue for each sample
-    ftype maxValue = -std::numeric_limits<ftype>::infinity();
-    for(tensorDim_t j=0; j<nCols; j++){
-      maxValue = std::max(maxValue, t.get(i, j));
-    }
+  auto res = t.createEmptyCopy();
 
-    for(tensorDim_t j=0; j<nCols; j++){
-      ftype e = t.get(i, j)-maxValue;
-      tmp.set(exp(e), i, j);
-    }
-  }
+  switch(t.getDevice()) {
+    case Device::CPU: {
+      const tensorSize_t stride = t.getDims()[-1];
 
-  const tensorSize_t stride = t.getDims()[-1];
-  Tensor res(t.getDims(), t.getDevice());
-  auto compute = [&res, &tmp, stride](tensorSize_t start){
-    ftype sum = 0;
-    for(tensorSize_t i=start; i<start+stride; i++){
-      sum += tmp[i];
-    }
+      // pre-compute exponents, centering each slice around its max for numerical stability
+      Tensor tmp(t.getDims(), t.getDevice(), false);
+      tensorSize_t offset = 0;
+      while(offset < t.getSize()) {
+        ftype maxValue = -std::numeric_limits<ftype>::infinity();
+        for(tensorSize_t i = offset; i < offset + stride; i++) {
+          maxValue = std::max(maxValue, t[i]);
+        }
 
-    for(tensorSize_t i=start; i<start+stride; i++){
-      res.set(tmp[i] / sum, i);
+        for(tensorSize_t i = offset; i < offset + stride; i++) {
+          tmp.set(exp(t[i] - maxValue), i);
+        }
+
+        offset += stride;
+      }
+
+      auto compute = [&res, &tmp, stride](tensorSize_t start){
+        ftype sum = 0;
+        for(tensorSize_t i = start; i < start + stride; i++){
+          sum += tmp[i];
+        }
+
+        for(tensorSize_t i = start; i < start + stride; i++){
+          res.set(tmp[i] / sum, i);
+        }
+      };
+
+      offset = 0;
+      while(offset < res.getSize()) {
+        compute(offset);
+        offset += stride;
+      }
+      break;
     }
-  };
-  
-  tensorSize_t offset=0;
-  while(offset<res.getSize()) {
-    compute(offset);
-    offset += stride;
+    case Device::CUDA:
+    #ifdef __CUDA
+      cuda_impl::softmax(res, t);
+    #else
+      __throw_invalid_argument("Attempted to give CUDA tensor");
+    #endif
+      break;
   }
 
   return res;
diff --git a/src/backend/module/layers/cuda/layers.cu b/src/backend/module/layers/cuda/layers.cu
new file mode 100644
index 0000000..1ba6c7b
--- /dev/null
+++ b/src/backend/module/layers/cuda/layers.cu
@@ -0,0 +1,34 @@
+/**
+ * @file layers.cu
+ * @author Robert Baumgartner (r.baumgartner-1@tudelft.nl)
+ * @brief
+ * @version 0.1
+ * @date 2026-05-13
+ *
+ * @copyright Copyright (c) 2026
+ *
+ */
+
+#ifndef __CUDA
+static_assert(false, "File should not be compiled without CUDA enabled");
+#endif // __CUDA
+
+#include "layers.cuh"
+#include "utility/cuda/cuda_common.cuh"
+
+using namespace std;
+
+namespace {
+  // TODO: matMulPlusBias kernel
+}
+
+namespace cuda_impl {
+  void matMulPlusBias(Tensor& res, const Tensor& input, const Tensor& weights, const Tensor& bias) {
+    constexpr int threadsPerBlock = 256;
+    const int blocks = (res.getSize() + threadsPerBlock - 1) / threadsPerBlock;
+
+    // TODO: launch kernel
+
+    cudaErrchk(cudaDeviceSynchronize());
+  }
+}
diff --git a/src/backend/module/layers/cuda/layers.cuh b/src/backend/module/layers/cuda/layers.cuh
new file mode 100644
index 0000000..1f74d78
--- /dev/null
+++ b/src/backend/module/layers/cuda/layers.cuh
@@ -0,0 +1,23 @@
+/**
+ * @file layers.cuh
+ * @author Robert Baumgartner (r.baumgartner-1@tudelft.nl)
+ * @brief
+ * @version 0.1
+ * @date 2026-05-13
+ *
+ * @copyright Copyright (c) 2026
+ *
+ */
+
+#pragma once
+
+#ifndef __CUDA
+static_assert(false, "File should not be included without CUDA enabled");
+#endif // __CUDA
+
+#include "utility/global_params.h"
+#include "data_modeling/tensor.h"
+
+namespace cuda_impl {
+  void matMulPlusBias(Tensor& res, const Tensor& input, const Tensor& weights, const Tensor& bias);
+}
diff --git a/src/backend/module/layers/ff_layer.cpp b/src/backend/module/layers/ff_layer.cpp
index 82fed69..2053ba5 100644
--- a/src/backend/module/layers/ff_layer.cpp
+++ b/src/backend/module/layers/ff_layer.cpp
@@ -17,6 +17,12 @@
 #include <cstdlib>
 #include <utility>
 
+#ifdef __CUDA
+#include "module/layers/cuda/layers.cuh"
+#else
+#include <stdexcept>
+#endif
+
 using namespace std;
 using namespace module;
 using namespace utility;
@@ -34,7 +40,7 @@ FfLayer::FfLayer(tensorDim_t inSize, tensorDim_t outSize, bool useBias, bool req
  */
 FfLayer::FfLayer(tensorDim_t inSize, tensorDim_t outSize, Device d, 
     bool useBias, bool requiresGrad, shared_ptr<InitializerBase> init)
-  : useBias{useBias}, requiresGrad{requiresGrad} 
+  : requiresGrad{requiresGrad} 
 {
   if(!init){
     init = make_shared<NormalXavierInitializer>(inSize, outSize);  
@@ -55,21 +61,47 @@ FfLayer::FfLayer(tensorDim_t inSize, tensorDim_t outSize, Device d,
  * Assumption for input: (b-size, ..., dim1, in-size)
  */
 Tensor FfLayer::operator()(const Tensor& input) const {
-  auto res = input.matmul(*weights);
+  switch(input.getDevice()) {
+    case Device::CPU: {
+      auto res = input.matmul(*weights);
+      if(bias) res = res + *bias;
+      return res;
+    }
+    case Device::CUDA:
+    #ifdef __CUDA
+      {
+        if(bias) {
+          // TODO: the following should be an optimized fusion kernel
+          /* auto resDims = input.getDims().toVector();
+          resDims.back() = static_cast<tensorDim_t>(weights->getDims().get(-1));
+          Tensor res(resDims, input.getDevice(), false);
 
-  if(useBias){
-    res = res + *bias;
+          cuda_impl::matMulPlusBias(res, input, *weights, *bias); */
+
+          auto res = input.matmul(*weights);
+          res = res + *bias;
+          return res;
+        } 
+
+        return input.matmul(*weights);
+      }
+    #else
+      __throw_invalid_argument("Attempted to give CUDA tensor");
+      return input.createShallowCopy(); // line should not be reached
+    #endif
   }
 
-  return res;
+  __throw_runtime_error("This line should never be reached.");
+  return input.createShallowCopy(); // suppress warnings
 }
 
 /**
  * @brief Like overload, but creates computational graph.
  */
 std::shared_ptr<Tensor> FfLayer::operator()(const std::shared_ptr<Tensor>& input) const {
+  // TODO: if you fuse kernel you'll also have to do it here. Perhaps give tensor a matmulplusbias method?
   auto res = cgraph::matmul(input, weights);
-  if(useBias){
+  if(bias){
     res = cgraph::add(res, bias);
   }
 
diff --git a/src/backend/module/layers/ff_layer.h b/src/backend/module/layers/ff_layer.h
index 2c99057..7ee5f50 100644
--- a/src/backend/module/layers/ff_layer.h
+++ b/src/backend/module/layers/ff_layer.h
@@ -19,7 +19,6 @@
 namespace module {
   class FfLayer : public ModuleBase {
     bool requiresGrad = false;
-    bool useBias = false;
 
     std::shared_ptr<Tensor> weights = nullptr;
     std::shared_ptr<Tensor> bias = nullptr;
@@ -50,6 +49,13 @@ namespace module {
       return {weights, bias};
     }
 
+    void setDevice(Device d) noexcept override {
+      weights->setDevice(d);
+      if(bias) {
+        bias->setDevice(d);
+      }
+    }
+
     void print(std::ostream& os) const noexcept override;
   };
 }
diff --git a/src/backend/module/module_base.cpp b/src/backend/module/module_base.cpp
index 951f96c..fdbcbbb 100644
--- a/src/backend/module/module_base.cpp
+++ b/src/backend/module/module_base.cpp
@@ -15,7 +15,9 @@
 
 using namespace std;
 
-ostream& module::operator<<(ostream& os, const module::ModuleBase& l) noexcept {
-  l.print(os); // calling vtable
-  return os;
+namespace module {
+  ostream& operator<<(ostream& os, const ModuleBase& l) noexcept {
+    l.print(os); // calling vtable
+    return os;
+  }
 }
\ No newline at end of file
diff --git a/src/backend/module/module_base.h b/src/backend/module/module_base.h
index 28247ab..6b05860 100644
--- a/src/backend/module/module_base.h
+++ b/src/backend/module/module_base.h
@@ -61,6 +61,8 @@ namespace module {
       #endif
       };
 
+      virtual void setDevice(Device d) noexcept {}
+
       friend std::ostream& operator<<(std::ostream& os, const ModuleBase& t) noexcept;
   };
 }
\ No newline at end of file
diff --git a/src/backend/shared/cuda/common_kernels.cuh b/src/backend/shared/cuda/common_kernels.cuh
new file mode 100644
index 0000000..1abf5e7
--- /dev/null
+++ b/src/backend/shared/cuda/common_kernels.cuh
@@ -0,0 +1,70 @@
+/**
+ * @file common_kernels.cuh
+ * @author Robert Baumgartner (r.baumgartner-1@tudelft.nl)
+ * @brief 
+ * @version 0.1
+ * @date 2026-05-25
+ * 
+ * @copyright Copyright (c) 2026
+ * 
+ */
+
+#pragma once
+
+#include "utility/global_params.h"
+#include "utility/utils.h"
+
+namespace cuda_impl {
+  template<typename T>
+  __device__ __forceinline__ ftype cudaMax(const ftype a, const ftype b) {
+    if constexpr (std::is_same_v<T, float>) {
+      return fmaxf(a, b);
+    } 
+    else if constexpr (std::is_same_v<T, double>) {
+        return fmax(a, b);
+    }
+    else {
+      static_assert(always_false<T>, "Unexpected value for ftype encountered");
+    }
+  }
+
+  /**
+    * @brief Single sigmoid computation.
+    */
+   template<typename T>
+  __device__ __forceinline__ ftype cudaSigmoid(const ftype x) {
+    if constexpr (std::is_same_v<T, float>) {
+      ftype z = expf(-fabsf(x));
+      ftype s = 1.0f / (1.0f + z);
+      return (x >= 0.f) ? s : z * s; // x < 0 => e^x/(e^x+1) 
+    } 
+    else if constexpr (std::is_same_v<T, double>) {
+      ftype z = exp(-abs(x));
+      ftype s = 1.0 / (1.0 + z);
+      return (x >= 0.0) ? s : z * s; // x < 0 => e^x/(e^x+1) 
+    }
+    else {
+      static_assert(always_false<T>, "Unexpected value for ftype encountered");
+    }
+  }
+
+  template<typename T>
+  __device__ __forceinline__ ftype cudaSqrt(const ftype x) {
+    if constexpr (std::is_same_v<T, float>) {
+      return sqrtf(x);
+    } 
+    else if constexpr (std::is_same_v<T, double>) {
+      return sqrt(x);
+    }
+    else {
+      static_assert(always_false<T>, "Unexpected value for ftype encountered");
+    }
+  }
+
+  /**
+   * @brief For single normalization, e.g. when normalizing with batch-size.
+   */
+  static __global__ void divideScalarKernel(ftype* const val, const ftype divisor) {
+    val[0] /= divisor;
+  }
+}
\ No newline at end of file
diff --git a/src/backend/shared/cuda/common_softmax.cuh b/src/backend/shared/cuda/common_softmax.cuh
new file mode 100644
index 0000000..9e3098d
--- /dev/null
+++ b/src/backend/shared/cuda/common_softmax.cuh
@@ -0,0 +1,226 @@
+/**
+ * @file common_softmax.cuh
+ * @author Robert Baumgartner (r.baumgartner-1@tudelft.nl)
+ * @brief Common kernels for softmax.
+ * @version 0.1
+ * @date 2026-05-25
+ * 
+ * @copyright Copyright (c) 2026
+ * 
+ */
+
+#pragma once
+
+#include "common_kernels.cuh"
+
+#include "utility/global_params.h"
+#include "utility/utils.h"
+
+namespace cuda_impl {
+
+  /**
+   * @brief Basic helper.
+   */
+  template<typename T>
+  __forceinline__ __device__ ftype stableExp(const ftype x, const ftype maxVal) {
+    if constexpr (std::is_same_v<T, float>) {
+      return __expf(x - maxVal);
+    }
+    else if constexpr (std::is_same_v<T, double>) {
+      return exp(x - maxVal);
+    }
+    else {
+      static_assert(always_false<T>, "ftype encountered unexpected type");
+    }
+  }
+
+  /**
+   * @brief Reduction kernel that computes the maximum within the size of 2 * warpsize at maximum.
+   */
+  template<int maxoffset>
+  __forceinline__ __device__ void warpMaxReduce(volatile ftype* const input, const tensorSize_t stride, const int offset) {
+    // TODO: warp shuffle for newer architectures
+    if(maxoffset == 32) {
+      if(offset + 32 < stride) input[offset] = cudaMax<ftype>(input[offset], input[offset + 32]);
+    }
+    if(maxoffset >= 16) {
+      if(offset + 16 < stride) input[offset] = cudaMax<ftype>(input[offset], input[offset + 16]);
+    }
+    if(maxoffset >= 8) {
+      if(offset + 8 < stride) input[offset] = cudaMax<ftype>(input[offset], input[offset + 8]);
+    }
+    if(maxoffset >= 4) {
+      if(offset + 4 < stride) input[offset] = cudaMax<ftype>(input[offset], input[offset + 4]);
+    }
+    if(maxoffset >= 2) {
+      if(offset + 2 < stride) input[offset] = cudaMax<ftype>(input[offset], input[offset + 2]);
+    }
+    if(maxoffset >= 1) {
+      if(offset + 1 < stride) input[offset] = cudaMax<ftype>(input[offset], input[offset + 1]);
+    }
+  }
+
+  /**
+   * @brief Here we find the maximum within 'stride'. Assumption: One warp does exactly one element of stride!
+   * Reduction via warp reduce. res has the maximum values stored.
+   */
+  template<int maxoffset>
+  static __global__ void findMaxKernelOneWarp(ftype* const res, const ftype* const input, const tensorSize_t stride, const tensorSize_t size) {
+    assert(blockDim.x % 32 == 0);
+
+    int gid = blockIdx.x * blockDim.x + threadIdx.x;
+    if(gid >= size)
+      return;
+
+    int tid = threadIdx.x;
+    extern __shared__ ftype smem[];
+    smem[tid] = input[gid];
+    __syncthreads();
+
+    volatile ftype* const start = smem + (tid / stride) * stride;
+    const int offset = gid % stride;
+    warpMaxReduce<maxoffset>(start, stride, offset);
+
+    // one warp reduces one 'stride'
+    if(offset == 0) {
+      res[tid / 32] = smem[tid];
+    }
+  }
+
+  /**
+   * @brief Like findMaxKernelOneWarp. The difference now is that the input size can be much larger. stride is 
+   * 64 < stride <= threadsPerBlock. res has the maximum values stored.
+   * 
+   * In this initial version we assume one kernel per stride, to make matters simple to understand.
+   */
+  static __global__ void findMaxKernelOneBlock(ftype* const res, const ftype* const input, const tensorSize_t stride) {
+    assert_debug(blockDim.x / stride == 0, "Kernel built for one stride per block, blockDim.x is < stride"); 
+
+    const int tid = threadIdx.x;
+    const int gid = blockIdx.x * stride + tid;
+
+    extern __shared__ ftype smem[]; // can lead to bank conflicts iff std::is_same_v<T, double>
+
+    const tensorSize_t maxIdx = tid + blockDim.x;
+    const bool doPadding = maxIdx >= stride;
+    if(doPadding) {
+      smem[tid] = input[gid];
+    }
+    else {
+      smem[tid] = cudaMax<ftype>(input[gid], input[gid + blockDim.x]);
+    }
+    __syncthreads();
+
+    for(tensorSize_t offset = blockDim.x >> 1; offset > 32; offset >>= 1) {
+      if(tid < offset) {
+        smem[tid] = cudaMax<ftype>(smem[tid], smem[tid + offset]);
+      }
+      __syncthreads();
+    }
+
+    // TODO: warp shuffle for newer architectures
+    volatile ftype* const start = smem;
+    if(tid < 32) {
+      start[tid] = cudaMax<ftype>(start[tid], start[tid + 32]);
+      start[tid] = cudaMax<ftype>(start[tid], start[tid + 16]);
+      start[tid] = cudaMax<ftype>(start[tid], start[tid + 8]);
+      start[tid] = cudaMax<ftype>(start[tid], start[tid + 4]);
+      start[tid] = cudaMax<ftype>(start[tid], start[tid + 2]);
+      start[tid] = cudaMax<ftype>(start[tid], start[tid + 1]);
+    }
+
+    if(tid == 0) { // one block per stride
+      res[blockIdx.x] = start[0];
+    }
+  }
+
+  /**
+   * @brief Like findMaxKernelOneBlock, but finding partial maximum. This is the case when the stride is too
+   * large to fit in one block.
+   */
+  static __global__ void findMaxKernelLargePass1(ftype* const partialMaxValues, const ftype* const input, 
+                                          const tensorSize_t stride, const int blocksPerStride) {
+    const int tid = threadIdx.x;
+    const int strideIdx = blockIdx.x / blocksPerStride;
+    const int blockWithinStride = blockIdx.x % blocksPerStride;
+
+    // block 0 within stride handles elements [0, 2*blockDim.x), block 1 within stride handles elements [2*blockDim.x, 4*blockDim.x), ...
+    const int inputBase = strideIdx * stride + blockWithinStride * 2 * blockDim.x; 
+      
+    extern __shared__ ftype smem[];
+    const tensorSize_t localIdx0 = inputBase + tid;
+    const tensorSize_t localIdx1 = inputBase + tid + blockDim.x;
+      
+    // localIdx0 < (strideIdx + 1) * stride <- checks whether thread idx exceeds bounds of this stride; one stride per block at cudaMax<ftype>
+    smem[tid] = (localIdx0 < (strideIdx + 1) * stride) ? input[localIdx0] : -INFINITY;
+    smem[tid + blockDim.x] = (localIdx1 < (strideIdx + 1) * stride) ? input[localIdx1] : -INFINITY;
+    __syncthreads();
+
+    // same reduction as findMaxKernelOneBlock from here
+    for(tensorSize_t offset = blockDim.x; offset > 32; offset >>= 1) {
+      if(tid < offset){
+        smem[tid] = cudaMax<ftype>(smem[tid], smem[tid + offset]);
+      } 
+      __syncthreads();
+    }
+
+    volatile ftype* start = smem;
+    if(tid < 32) {
+      start[tid] = cudaMax<ftype>(start[tid], start[tid + 32]);
+      start[tid] = cudaMax<ftype>(start[tid], start[tid + 16]);
+      start[tid] = cudaMax<ftype>(start[tid], start[tid + 8]);
+      start[tid] = cudaMax<ftype>(start[tid], start[tid + 4]);
+      start[tid] = cudaMax<ftype>(start[tid], start[tid + 2]);
+      start[tid] = cudaMax<ftype>(start[tid], start[tid + 1]);
+    }
+
+    if(tid == 0) {
+      partialMaxValues[blockIdx.x] = start[0];
+    }
+  }
+
+  /**
+   * @brief Self explanatory following findMaxKernelLargePass1. Assumption: All remaining cudaMax<ftype> values do fit into 
+   * one single block now -> we launch one block per stride this time.
+   */
+  static __global__ void findMaxKernelLargePass2(ftype* const maxValues, const ftype* const partialMaxValues, const tensorSize_t blocksPerStride) {
+    assert_debug(blockDim.x / blocksPerStride == 0, "Kernel built for one stride per block, blockDim.x is < stride"); 
+
+    const int tid = threadIdx.x;
+    const int gid = blockIdx.x * blocksPerStride + tid;
+
+    extern __shared__ ftype smem[]; // can lead to bank conflicts iff std::is_same_v<T, double>
+
+    const tensorSize_t maxIdx = tid + blockDim.x;
+    const bool doPadding = maxIdx >= blocksPerStride;
+    if(doPadding) {
+      smem[tid] = partialMaxValues[gid];
+    }
+    else {
+      smem[tid] = cudaMax<ftype>(partialMaxValues[gid], partialMaxValues[gid + blockDim.x]);
+    }
+    __syncthreads();
+
+    for(tensorSize_t offset = blockDim.x >> 1; offset > 32; offset >>= 1) {
+      if(tid < offset) {
+        smem[tid] = cudaMax<ftype>(smem[tid], smem[tid + offset]);
+      }
+      __syncthreads();
+    }
+
+    // TODO: warp shuffle for newer architectures
+    volatile ftype* const start = smem;
+    if(tid < 32) {
+      if(tid + 32 < blockDim.x) start[tid] = cudaMax<ftype>(start[tid], start[tid + 32]);
+      if(tid + 16 < blockDim.x) start[tid] = cudaMax<ftype>(start[tid], start[tid + 16]);
+      if(tid + 8  < blockDim.x) start[tid] = cudaMax<ftype>(start[tid], start[tid + 8]);
+      if(tid + 4  < blockDim.x) start[tid] = cudaMax<ftype>(start[tid], start[tid + 4]);
+      if(tid + 2  < blockDim.x) start[tid] = cudaMax<ftype>(start[tid], start[tid + 2]);
+      if(tid + 1  < blockDim.x) start[tid] = cudaMax<ftype>(start[tid], start[tid + 1]);
+    }
+
+    if(tid == 0) { // one block per stride
+      maxValues[blockIdx.x] = start[0];
+    }
+  }
+}
\ No newline at end of file
diff --git a/src/backend/training/loss_functions/bce_loss.cpp b/src/backend/training/loss_functions/bce_loss.cpp
index a6df909..8176784 100644
--- a/src/backend/training/loss_functions/bce_loss.cpp
+++ b/src/backend/training/loss_functions/bce_loss.cpp
@@ -1,20 +1,25 @@
 /**
  * @file bce_loss.cpp
  * @author Robert Baumgartner (r.baumgartner-1@tudelft.nl)
- * @brief 
+ * @brief
  * @version 0.1
  * @date 2026-03-07
- * 
+ *
  * @copyright Copyright (c) 2026
- * 
+ *
  */
 
 #include "bce_loss.h"
-
 #include "computational_graph/loss_functions/bce_node.h"
 
 #include <cmath>
 
+#ifdef __CUDA
+#include "training/loss_functions/cuda/loss_functions.cuh"
+#else
+#include <stdexcept>
+#endif
+
 using namespace std;
 using namespace train;
 
@@ -25,7 +30,7 @@ using namespace train;
 shared_ptr<Tensor> BceLoss::operator()(const shared_ptr<Tensor> y, const shared_ptr<Tensor> ypred) const {
   if(!ypred->getRequiresGrad()) {
     __throw_invalid_argument("ypred must have gradient enabled");
-  }  
+  }
   else if(y->getDevice() != ypred->getDevice()){
     __throw_invalid_argument("y and ypred must be on same device");
   }
@@ -33,20 +38,35 @@ shared_ptr<Tensor> BceLoss::operator()(const shared_ptr<Tensor> y, const shared_
     __throw_invalid_argument("Tensors must be of same shape");
   }
 
-  auto bce = [](ftype y, ftype ypred){
-    return y*log(std::max(ypred, epsBce)) + (1-y)*log(std::max(1-ypred, epsBce));
-  };
+  shared_ptr<Tensor> res = nullptr;
 
-  const auto nBatches = y->getDims()[0];
+  switch(y->getDevice()) {
+    case Device::CPU: 
+    {
+      auto bce = [](const ftype y, const ftype ypred){
+        return y * log(std::max(ypred, EPS_BCE)) + (1 - y) * log(std::max(1 - ypred, EPS_BCE));
+      };
 
-  ftype loss = 0;
-  for(tensorSize_t i=0; i<nBatches; i++){
-    loss += bce((*y)[i], (*ypred)[i]);
+      const auto nBatches = y->getDims()[0];
+      ftype loss = 0;
+      for(tensorSize_t i = 0; i < nBatches; i++){
+        loss += bce((*y)[i], (*ypred)[i]);
+      }
+      
+      res = make_shared<Tensor>(std::vector<tensorDim_t>{1}, std::vector<ftype>{-loss / nBatches}, Device::CPU, true);
+      break;
+    }
+    case Device::CUDA:
+    #ifdef __CUDA
+      res = make_shared<Tensor>(vector<tensorDim_t>{1}, Device::CUDA, true);
+      cuda_impl::bceLoss(*res, *y, *ypred);
+    #else
+      __throw_invalid_argument("Attempted to give CUDA tensor");
+    #endif
+      break;
   }
 
-  auto res = make_shared<Tensor>(std::vector<tensorDim_t>{1}, std::vector<ftype>{-loss / nBatches}, y->getDevice(), true);
   res->setCgNode(make_shared<cgraph::BceNode>(y, ypred));
   assert(res->getRequiresGrad());
-
-  return res; 
-}
\ No newline at end of file
+  return res;
+}
diff --git a/src/backend/training/loss_functions/bce_sigmoid_loss.cpp b/src/backend/training/loss_functions/bce_sigmoid_loss.cpp
index 2634bf4..eafbc05 100644
--- a/src/backend/training/loss_functions/bce_sigmoid_loss.cpp
+++ b/src/backend/training/loss_functions/bce_sigmoid_loss.cpp
@@ -1,19 +1,24 @@
 /**
  * @file bce_logits_loss.cpp
  * @author Robert Baumgartner (r.baumgartner-1@tudelft.nl)
- * @brief 
+ * @brief
  * @version 0.1
  * @date 2026-03-17
- * 
+ *
  * @copyright Copyright (c) 2026
- * 
+ *
  */
 
- #include "bce_sigmoid_loss.h"
+#include "bce_sigmoid_loss.h"
+#include "computational_graph/loss_functions/bce_sigmoid_node.h"
 
- #include "computational_graph/loss_functions/bce_sigmoid_node.h"
+#include <cmath>
 
- #include <cmath>
+#ifdef __CUDA
+#include "training/loss_functions/cuda/loss_functions.cuh"
+#else
+#include <stdexcept>
+#endif
 
 using namespace std;
 using namespace train;
@@ -25,7 +30,7 @@ using namespace train;
 shared_ptr<Tensor> BceSigmoidLoss::operator()(const shared_ptr<Tensor> y, const shared_ptr<Tensor> logits) const {
   if(!logits->getRequiresGrad()) {
     __throw_invalid_argument("logits must have gradient enabled");
-  }  
+  }
   else if(y->getDevice() != logits->getDevice()){
     __throw_invalid_argument("y and logits must be on same device");
   }
@@ -33,21 +38,34 @@ shared_ptr<Tensor> BceSigmoidLoss::operator()(const shared_ptr<Tensor> y, const
     __throw_invalid_argument("Tensors must be of same shape");
   }
 
-  auto bceSimplified = [](ftype y, ftype logit){
-    constexpr ftype zero = 0;
-    return std::max(logit, zero) - logit*y + log(1+exp(-std::abs(logit)));
-  };
+  shared_ptr<Tensor> res = nullptr;
 
-  const auto nBatches = y->getDims()[0];
+  switch(y->getDevice()) {
+    case Device::CPU: {
+      auto bceSimplified = [](ftype y, ftype logit){
+        constexpr ftype zero = 0;
+        return std::max(logit, zero) - logit * y + log(1 + exp(-std::abs(logit)));
+      };
 
-  ftype loss = 0;
-  for(tensorSize_t i=0; i<nBatches; i++){
-    loss += bceSimplified((*y)[i], (*logits)[i]);
+      const auto nBatches = y->getDims()[0];
+      ftype loss = 0;
+      for(tensorSize_t i = 0; i < nBatches; i++){
+        loss += bceSimplified((*y)[i], (*logits)[i]);
+      }
+      res = make_shared<Tensor>(std::vector<tensorDim_t>{1}, std::vector<ftype>{loss / nBatches}, y->getDevice(), true);
+      break;
+    }
+    case Device::CUDA:
+    #ifdef __CUDA
+      res = make_shared<Tensor>(vector<tensorDim_t>{1}, Device::CUDA, true);
+      cuda_impl::bceSigmoidLoss(*res, *y, *logits);
+    #else
+      __throw_invalid_argument("Attempted to give CUDA tensor");
+    #endif
+      break;
   }
 
-  auto res = make_shared<Tensor>(std::vector<tensorDim_t>{1}, std::vector<ftype>{loss / nBatches}, y->getDevice(), true);
   res->setCgNode(make_shared<cgraph::BceSigmoidNode>(y, logits));
   assert(res->getRequiresGrad());
-
-  return res; 
-}
\ No newline at end of file
+  return res;
+}
diff --git a/src/backend/training/loss_functions/crossentropy_loss.cpp b/src/backend/training/loss_functions/crossentropy_loss.cpp
index d1a5291..641ab3e 100644
--- a/src/backend/training/loss_functions/crossentropy_loss.cpp
+++ b/src/backend/training/loss_functions/crossentropy_loss.cpp
@@ -1,20 +1,25 @@
 /**
  * @file crossentropy_loss.cpp
  * @author Robert Baumgartner (r.baumgartner-1@tudelft.nl)
- * @brief 
+ * @brief
  * @version 0.1
  * @date 2026-03-17
- * 
+ *
  * @copyright Copyright (c) 2026
- * 
+ *
  */
 
 #include "crossentropy_loss.h"
-
 #include "computational_graph/loss_functions/crossentropy_node.h"
 
 #include <cmath>
 
+#ifdef __CUDA
+#include "training/loss_functions/cuda/loss_functions.cuh"
+#else
+#include <stdexcept>
+#endif
+
 using namespace std;
 using namespace train;
 
@@ -33,23 +38,32 @@ shared_ptr<Tensor> CrossEntropyLoss::operator()(const shared_ptr<Tensor> y, cons
     __throw_invalid_argument("Tensors must be of same shape");
   }
 
-  auto ce = [&y, &ypred](const tensorDim_t b){
-    ftype res = 0;
-    for(tensorDim_t i=0; i<y->getDims()[-1]; i++){
-      res += y->get(b, i) * log(std::max(ypred->get(b, i), epsCrossentropy));
-    }
-    return res;
-  };
+  shared_ptr<Tensor> res = nullptr;
+
+  switch(y->getDevice()) {
+    case Device::CPU: {
+      const tensorSize_t stride = y->getDims()[-1];
+      const tensorSize_t nSamples = y->getSize() / stride;
 
-  const auto nBatches = y->getDims()[0];
-  ftype loss = 0;
-  for(tensorSize_t b=0; b<nBatches; b++){
-    loss += ce(b);
+      ftype loss = 0;
+      for (tensorSize_t i = 0; i < y->getSize(); i++) {
+        loss += (*y)[i] * log(std::max((*ypred)[i], EPS_CROSSENTROPY));
+      }
+
+      res = make_shared<Tensor>(std::vector<tensorDim_t>{1}, std::vector<ftype>{-loss / nSamples}, y->getDevice(), true);
+      break;
+    }
+    case Device::CUDA:
+    #ifdef __CUDA
+      res = make_shared<Tensor>(vector<tensorDim_t>{1}, Device::CUDA, true);
+      cuda_impl::crossEntropyLoss(*res, *y, *ypred);
+    #else
+      __throw_invalid_argument("Attempted to give CUDA tensor");
+    #endif
+      break;
   }
 
-  auto res = make_shared<Tensor>(std::vector<tensorDim_t>{1}, std::vector<ftype>{-loss / nBatches}, y->getDevice(), true);
   res->setCgNode(std::make_shared<cgraph::CrossEntropyNode>(y, ypred));
   assert(res->getRequiresGrad());
-
   return res;
-}
\ No newline at end of file
+}
diff --git a/src/backend/training/loss_functions/crossentropy_softmax_loss.cpp b/src/backend/training/loss_functions/crossentropy_softmax_loss.cpp
index a2b7866..feedf8d 100644
--- a/src/backend/training/loss_functions/crossentropy_softmax_loss.cpp
+++ b/src/backend/training/loss_functions/crossentropy_softmax_loss.cpp
@@ -1,20 +1,25 @@
 /**
  * @file crossentropy_logits_loss.cpp
- * @author Robert Baumgartner (r.baumgartner-1@tudelflogits->nl)
- * @brief 
+ * @author Robert Baumgartner (r.baumgartner-1@tudelft.nl)
+ * @brief
  * @version 0.1
  * @date 2026-03-17
- * 
+ *
  * @copyright Copyright (c) 2026
- * 
+ *
  */
 
 #include "crossentropy_softmax_loss.h"
-
 #include "computational_graph/loss_functions/crossentropy_softmax_node.h"
 
 #include <cmath>
 
+#ifdef __CUDA
+#include "training/loss_functions/cuda/loss_functions.cuh"
+#else
+#include <stdexcept>
+#endif
+
 using namespace std;
 using namespace train;
 
@@ -33,61 +38,73 @@ shared_ptr<Tensor> CrossEntropySoftmaxLoss::operator()(const shared_ptr<Tensor>
     __throw_invalid_argument("Tensors must be of same shape");
   }
 
-  ////////////////////////////////////////////////
+  shared_ptr<Tensor> res = nullptr;
 
-  const auto nRows = logits->getDims()[-2];
-  const auto nCols = logits->getDims()[-1];
+  switch(y->getDevice()) {
+    case Device::CPU: {
+      const tensorSize_t stride = logits->getDims()[-1];
+      const tensorSize_t nSamples = logits->getSize() / stride;
 
-  // pre-compute exponents and max-values
-  vector<ftype> maxValues(nRows); 
-  Tensor tmp(logits->getDims(), logits->getDevice(), false);
-  for(tensorDim_t i=0; i<nRows; i++){
-    // for numerical stability, avoid inf
-    ftype maxV = -std::numeric_limits<ftype>::infinity();
-    for(tensorDim_t j=0; j<nCols; j++){
-      maxV = std::max(maxV, logits->get(i, j));
-    }
+      // pre-compute exponents and max-values, centering each slice for numerical stability
+      vector<ftype> maxValues(nSamples);
+      Tensor tmp(logits->getDims(), logits->getDevice(), false);
+      tensorSize_t offset = 0;
+      while(offset < logits->getSize()) {
+        ftype maxV = -std::numeric_limits<ftype>::infinity();
+        for(tensorSize_t i = offset; i < offset + stride; i++) {
+          maxV = std::max(maxV, (*logits)[i]);
+        }
 
-    maxValues[i] = maxV;
+        maxValues[offset / stride] = maxV;
 
-    for(tensorDim_t j=0; j<nCols; j++){
-      ftype e = logits->get(i, j)-maxV;
-      tmp.set(exp(e), i, j);
-    }
-  }
+        for(tensorSize_t i = offset; i < offset + stride; i++) {
+          tmp.set(exp((*logits)[i] - maxV), i);
+        }
 
-  const tensorSize_t stride = logits->getDims()[-1];
-  ftype loss = 0;
-
-  /** 
-   * CE = -sum_i(y_i * z_i) + log(sum_j(exp(z_j))) with 
-   * log(sum_j(exp(z_j))) = max(z) + log(sum_j(exp(z_j - max(z)))).
-   * for numerical stability 
-   */ 
-  auto compute = [&loss, &y, &logits, &tmp, &maxValues, stride](tensorSize_t start){
-    ftype lsum = 0;
-    for(tensorSize_t i=start; i<start+stride; i++){
-      lsum += tmp[i];
-    }
-    lsum = log(lsum);
+        offset += stride;
+      }
+      ftype loss = 0;
+
+      /**
+       * CE = -sum_i(y_i * z_i) + log(sum_j(exp(z_j))) with
+       * log(sum_j(exp(z_j))) = max(z) + log(sum_j(exp(z_j - max(z)))).
+       * for numerical stability
+       */
+      auto compute = [&loss, &y, &logits, &tmp, &maxValues, stride](tensorSize_t start) {
+        ftype lsum = 0;
+        for(tensorSize_t i = start; i < start + stride; i++){
+          lsum += tmp[i];
+        }
+        lsum = log(lsum);
+
+        const tensorSize_t j = start / stride;
+        for(tensorSize_t i = start; i < start + stride; i++) {
+          if((*y)[i]>0){ // y either zero or one
+            loss += -(*logits)[i] + maxValues[j] + lsum;
+          }
+        }
+      };
 
-    const tensorSize_t j = start/stride;
-    for(tensorSize_t i=start; i<start+stride; i++){
-      if((*y)[i]>0){ // y either zero or one
-        loss += -(*logits)[i] + maxValues[j] + lsum;
+      offset = 0;
+      while(offset < logits->getSize()) {
+        compute(offset);
+        offset += stride;
       }
+
+      res = make_shared<Tensor>(std::vector<tensorDim_t>{1}, std::vector<ftype>{loss / static_cast<ftype>(nSamples)}, y->getDevice(), true);
+      break;
     }
-  };
-  
-  tensorSize_t offset=0;
-  while(offset<logits->getSize()) {
-    compute(offset);
-    offset += stride;
+    case Device::CUDA:
+    #ifdef __CUDA
+      res = make_shared<Tensor>(vector<tensorDim_t>{1}, Device::CUDA, true);
+      cuda_impl::crossEntropySoftmaxLoss(*res, *y, *logits);
+    #else
+      __throw_invalid_argument("Attempted to give CUDA tensor");
+    #endif
+      break;
   }
 
-  auto res = make_shared<Tensor>(std::vector<tensorDim_t>{1}, std::vector<ftype>{loss / logits->getDims()[0]}, y->getDevice(), true);
   res->setCgNode(std::make_shared<cgraph::CrossEntropySoftmaxNode>(y, logits));
   assert(res->getRequiresGrad());
-
   return res;
-}
\ No newline at end of file
+}
diff --git a/src/backend/training/loss_functions/cuda/loss_functions.cu b/src/backend/training/loss_functions/cuda/loss_functions.cu
new file mode 100644
index 0000000..717e5b6
--- /dev/null
+++ b/src/backend/training/loss_functions/cuda/loss_functions.cu
@@ -0,0 +1,539 @@
+/**
+ * @file loss_functions.cu
+ * @author Robert Baumgartner (r.baumgartner-1@tudelft.nl)
+ * @brief
+ * @version 0.1
+ * @date 2026-05-10
+ *
+ * @copyright Copyright (c) 2026
+ *
+ */
+
+#ifndef __CUDA
+static_assert(false, "File should not be compiled without CUDA enabled");
+#endif // __CUDA
+
+#include "loss_functions.cuh"
+
+#include "utility/utils.h"
+#include "utility/cuda/cuda_common.cuh"
+
+#include "shared/cuda/common_kernels.cuh"
+#include "shared/cuda/common_softmax.cuh"
+
+#include <thrust/device_ptr.h>
+#include <thrust/reduce.h>
+
+using namespace std;
+
+namespace {
+  using namespace cuda_impl;
+
+  template<typename T>
+  __forceinline__ __device__ T bce(T y, T ypred) {
+    if constexpr (std::is_same_v<T, float>) {
+      return y * __logf(cudaMax<T>(ypred, EPS_BCE)) + (1 - y) * __logf(cudaMax<T>(1 - ypred, EPS_BCE));
+    }
+    else if constexpr (std::is_same_v<T, double>) {
+      return y * log(cudaMax<T>(ypred, EPS_BCE)) + (1 - y) * log(cudaMax<T>(1 - ypred, EPS_BCE));
+    }
+    else {
+      static_assert(always_false<T>, "Unexpected value for ftype");
+    }
+  }
+
+  /**
+   * @brief Forward BCE loss.
+   */
+  __global__ void bceLossKernel(ftype* const res, const ftype* const y, const ftype* const ypred, tensorSize_t size) {
+    const int gid = blockDim.x * blockIdx.x + threadIdx.x;
+    const int tid = threadIdx.x;
+
+    extern __shared__ ftype smem[];
+    const ftype tmp = gid < size ? bce<ftype>(y[gid], ypred[gid]) : 0;
+    smem[tid] = tmp;
+    __syncthreads();
+
+    for(tensorSize_t offset = blockDim.x / 2; offset > 32; offset >>= 1){
+      if(tid < offset) {
+        smem[tid] += smem[tid + offset];
+      }
+      __syncthreads();
+    }
+
+    // TODO: warp shuffles
+    volatile ftype* sdata = smem;
+    if(tid < 32) {
+      if(gid + 32 < size) {
+        sdata[tid] += sdata[tid + 32];
+      }
+      if(gid + 16 < size) {
+        sdata[tid] += sdata[tid + 16];
+      }
+      if(gid + 8 < size) {
+        sdata[tid] += sdata[tid + 8];
+      }
+      if(gid + 4 < size) {
+        sdata[tid] += sdata[tid + 4];
+      }
+      if(gid + 2 < size) {
+        sdata[tid] += sdata[tid + 2];
+      }
+      if(gid + 1 < size) {
+        sdata[tid] += sdata[tid + 1];
+      }
+    }
+
+    if(tid == 0) {
+      res[blockIdx.x] = sdata[0];
+    }
+  }
+
+  template<typename T>
+  __forceinline__ __device__ T bceSimplified(T y, T logit) {
+    constexpr T zero = 0;
+    if constexpr (std::is_same_v<T, float>) {
+      return cudaMax<T>(logit, zero) - logit * y + __logf(1 + __expf(-abs(logit)));
+    }
+    else if constexpr (std::is_same_v<T, double>) {
+      return cudaMax<T>(logit, zero) - logit * y + log(1 + exp(-abs(logit)));
+    }
+    else {
+      static_assert(always_false<T>, "Unexpected value for ftype");
+    }
+  }
+
+  /**
+   * @brief BCE kernel with integrated sigmoid.
+   * 
+   * @param logits Forward logits.
+   */
+  __global__ void bceSigmoidLossKernel(ftype* const res, const ftype* const y, const ftype* const logits, tensorSize_t size) {
+    const int gid = blockDim.x * blockIdx.x + threadIdx.x;    
+    const int tid = threadIdx.x;
+
+    extern __shared__ ftype smem[];
+
+    // pre-load first round
+    const ftype tmp = gid < size ? bceSimplified<ftype>(y[gid], logits[gid]) : 0;
+    smem[tid] = tmp;
+    __syncthreads();
+
+    for(tensorSize_t offset = blockDim.x / 2; offset > 32; offset >>= 1){
+      if(tid < offset) {
+        smem[tid] += smem[tid + offset];
+      }
+      __syncthreads();
+    }
+
+    // TODO: warp shuffles
+    volatile ftype* sdata = smem;
+    if(tid < 32) {
+      if(gid + 32 < size) {
+        sdata[tid] += sdata[tid + 32];
+      }
+      if(gid + 16 < size) {
+        sdata[tid] += sdata[tid + 16];
+      }
+      if(gid + 8 < size) {
+        sdata[tid] += sdata[tid + 8];
+      }
+      if(gid + 4 < size) {
+        sdata[tid] += sdata[tid + 4];
+      }
+      if(gid + 2 < size) {
+        sdata[tid] += sdata[tid + 2];
+      }
+      if(gid + 1 < size) {
+        sdata[tid] += sdata[tid + 1];
+      }
+    }
+
+    if(tid == 0) {
+      res[blockIdx.x] = sdata[0];
+    }
+  }
+
+  template<typename T>
+  __forceinline__ __device__ ftype crossEntropy(const ftype y, const ftype ypred) {
+    if constexpr (std::is_same_v<T, float>) {
+      return y * __logf(cudaMax<T>(ypred, EPS_CROSSENTROPY));
+    }
+    else if constexpr (std::is_same_v<T, double>) {
+      return y * log(cudaMax<T>(ypred, EPS_CROSSENTROPY));
+    }
+    else {
+      static_assert(always_false<T>, "Encountered unexpected ftype");
+    }
+  }
+
+  /**
+   * @brief Cross-entropy reduction over a full block. Covers two times blockDim.x.
+   */
+  __global__ void crossEntropyLossKernelOneBlock(ftype* const res, const ftype* const y, const ftype* const yPred, const tensorSize_t size) {
+    const int tid = threadIdx.x;
+    const int gid = blockIdx.x * blockDim.x + tid;
+
+    extern __shared__ ftype smem[];
+
+    const ftype tmp = gid < size ? crossEntropy<ftype>(y[gid], yPred[gid]) : 0;
+    smem[tid] = tmp;
+    __syncthreads();
+
+    for(int offset = blockDim.x / 2; offset > 32; offset >>= 1) {
+      if(tid < offset) {
+        smem[tid] += smem[tid + offset];
+      }
+      __syncthreads();
+    }
+
+    // TODO: warp shuffle again
+    volatile ftype* sdata = smem;
+    if(tid < 32) {
+      if(gid + 32 < size) {
+        sdata[tid] += sdata[tid + 32];
+      }
+      if(gid + 16 < size) {
+        sdata[tid] += sdata[tid + 16];
+      }
+      if(gid + 8 < size) {
+        sdata[tid] += sdata[tid + 8];
+      }
+      if(gid + 4 < size) {
+        sdata[tid] += sdata[tid + 4];
+      }
+      if(gid + 2 < size) {
+        sdata[tid] += sdata[tid + 2];
+      }
+      if(gid + 1 < size) {
+        sdata[tid] += sdata[tid + 1];
+      }
+    }
+
+    if(threadIdx.x == 0) {
+      res[blockIdx.x] = sdata[0];
+    }
+  }
+
+  /**
+   * @brief Softmax cross-entropy loss kernel. One block per sample.
+   */
+
+  template<typename T>
+  __global__ void crossEntropySoftmaxLossKernel(ftype* const perSampleLoss, const ftype* const y, const ftype* const logits,
+                                                const ftype* const maxValues, const tensorSize_t stride) {
+    const int tid  = threadIdx.x;
+    const int tid2 = tid + blockDim.x;
+    const int sampleBase = blockIdx.x * stride;
+    const ftype maxV = maxValues[blockIdx.x];
+
+    extern __shared__ ftype smem[];
+
+    const bool inBounds0 = (tid < stride);
+    const bool inBounds1 = (tid2 < stride);
+
+    constexpr ftype zero = 0;
+    ftype e0  = inBounds0 ? stableExp<T>(logits[sampleBase + tid],  maxV) : zero;
+    ftype yz0 = inBounds0 ? y[sampleBase + tid]  * logits[sampleBase + tid]  : zero;
+    ftype e1  = inBounds1 ? stableExp<T>(logits[sampleBase + tid2], maxV) : zero;
+    ftype yz1 = inBounds1 ? y[sampleBase + tid2] * logits[sampleBase + tid2] : zero;
+
+    // Phase 1: reduce exp sum across stride (smem has 2 * blockDim.x slots)
+    smem[tid]  = e0;
+    smem[tid2] = e1;
+    __syncthreads();
+
+    for(int offset = blockDim.x; offset >= 1; offset >>= 1) {
+      if(tid < offset) {
+        smem[tid] += smem[tid + offset];
+      }
+      __syncthreads();
+    }
+
+    const ftype expSum = smem[0];
+
+    // reduce y*z dot product (y is one-hot, so at most one non-zero per sample)
+    smem[tid] = yz0 + yz1;
+    __syncthreads();
+
+    for(int offset = blockDim.x / 2; offset >= 1; offset >>= 1) {
+      if(tid < offset) {
+        smem[tid] += smem[tid + offset];
+      }
+      __syncthreads();
+    }
+
+    if(tid == 0) {
+      ftype logExpSum;
+      if constexpr (std::is_same_v<T, float>) {
+        logExpSum = __logf(expSum);
+      }
+      else if constexpr (std::is_same_v<T, double>) {
+        logExpSum = log(expSum);
+      }
+      else {
+        static_assert(always_false<T>, "Unexpected value for ftype encountered");
+      }
+      perSampleLoss[blockIdx.x] = maxV + logExpSum - smem[0];
+    }
+  }
+
+  /**
+   * @brief Helper for RMSE loss.
+   */
+  __forceinline__ __device__ ftype diffPow(const ftype y, const ftype ypred) {
+    auto diff = y - ypred;
+    return diff * diff;
+  }
+
+  /**
+   * @brief RMSE forward loss.
+   */
+  __global__ void rmseKernelOneBlock(ftype* const res, const ftype* const y, const ftype* const yPred, const tensorSize_t size) {
+    const int tid = threadIdx.x;
+    const int gid = blockIdx.x * blockDim.x + tid;
+
+    extern __shared__ ftype smem[];
+    const ftype tmp = gid < size ? diffPow(y[gid], yPred[gid]) : 0;
+    smem[tid] = tmp;
+    __syncthreads();
+
+    for(int offset = blockDim.x / 2; offset > 32; offset >>= 1) {
+      if(tid < offset) {
+        smem[tid] += smem[tid + offset];
+      }
+      __syncthreads();
+    }
+
+    // TODO: warp shuffle again
+    volatile ftype* sdata = smem;
+    if(tid < 32) {
+      if(gid + 32 < size) {
+        sdata[tid] += sdata[tid + 32];
+      }
+      if(gid + 16 < size) {
+        sdata[tid] += sdata[tid + 16];
+      }
+      if(gid + 8 < size) {
+        sdata[tid] += sdata[tid + 8];
+      }
+      if(gid + 4 < size) {
+        sdata[tid] += sdata[tid + 4];
+      }
+      if(gid + 2 < size) {
+        sdata[tid] += sdata[tid + 2];
+      }
+      if(gid + 1 < size) {
+        sdata[tid] += sdata[tid + 1];
+      }
+    }
+
+    if(threadIdx.x == 0) {
+      res[blockIdx.x] = sdata[0];
+    }
+  }
+
+  template<typename T>
+  __global__ void normalizeRmse(ftype* const val, ftype divisor) {
+    const ftype v = val[0];
+    if constexpr (std::is_same_v<T, float>) {
+      val[0] = sqrtf(v / divisor);
+    }
+    else if constexpr (std::is_same_v<T, double>) {
+      val[0] = sqrt(v / divisor);
+    }
+    else {
+      static_assert(always_false<T>, "Encountered unexpected ftype");
+    }
+  }
+}
+
+namespace cuda_impl {
+  void bceLoss(Tensor& res, const Tensor& y, const Tensor& yPred) {
+    constexpr int threadsPerBlock = 256;
+    const int blocks = (y.getSize() + threadsPerBlock - 1) / threadsPerBlock;
+
+    if(blocks > 1) {
+      // two pass solution
+      ftype* tmp; // TODO: Keep this guy in memory for an instance
+      cudaErrchk(cudaMalloc(&tmp, blocks * sizeof(ftype)));
+
+      bceLossKernel<<<blocks, threadsPerBlock, threadsPerBlock * sizeof(ftype)>>>(
+          tmp, y.getData(), yPred.getData(), y.getSize());
+      cudaErrchk(cudaDeviceSynchronize());
+
+      // do a sum over the residual array
+      thrust::device_ptr<ftype> tmpPtr(tmp);
+      thrust::device_ptr<ftype> resPtr(res.getData());
+      resPtr[0] = thrust::reduce(tmpPtr, tmpPtr + blocks, static_cast<ftype>(0.0f), thrust::plus<ftype>());
+
+      cudaErrchk(cudaFree(tmp));
+    }
+    else {
+      bceLossKernel<<<blocks, threadsPerBlock, threadsPerBlock * sizeof(ftype)>>>(
+          res.getData(), y.getData(), yPred.getData(), y.getSize());
+      cudaErrchk(cudaDeviceSynchronize());
+    }
+
+    // loss = -loss / nBatches
+    divideScalarKernel<<<1, 1>>>(res.getData(), -1 * static_cast<ftype>(y.getDims()[0]));
+    cudaErrchk(cudaDeviceSynchronize());
+  }
+
+  void bceSigmoidLoss(Tensor& res, const Tensor& y, const Tensor& logits) {
+    constexpr int threadsPerBlock = 256;
+    const int blocks = (y.getSize() + threadsPerBlock - 1) / threadsPerBlock;
+
+    if(blocks > 1) {
+      // we do two passes at max
+      ftype* tmp; 
+      cudaErrchk(cudaMalloc(&tmp, blocks * sizeof(ftype)));
+
+      bceSigmoidLossKernel<<<blocks, threadsPerBlock, threadsPerBlock * sizeof(ftype)>>>(
+          tmp, y.getData(), logits.getData(), y.getSize());
+      cudaErrchk(cudaDeviceSynchronize());
+
+      // do a sum over the residual array
+      thrust::device_ptr<ftype> tmpPtr(tmp);
+      thrust::device_ptr<ftype> resPtr(res.getData());
+      resPtr[0] = thrust::reduce(tmpPtr, tmpPtr + blocks, static_cast<ftype>(0.0f), thrust::plus<ftype>());
+
+      cudaErrchk(cudaFree(tmp));
+    }
+    else {
+      bceSigmoidLossKernel<<<blocks, threadsPerBlock, threadsPerBlock * sizeof(ftype)>>>(
+          res.getData(), y.getData(), logits.getData(), y.getSize());
+      cudaErrchk(cudaDeviceSynchronize());
+    }
+
+    // loss = loss / nBatches
+    divideScalarKernel<<<1, 1>>>(res.getData(), static_cast<ftype>(y.getDims()[0]));
+    cudaErrchk(cudaDeviceSynchronize());
+  }
+
+  void crossEntropyLoss(Tensor& res, const Tensor& y, const Tensor& yPred) {    
+    if(y.getSize() <= 256) {
+      int threadsPerBlock = 1;
+      while(threadsPerBlock < y.getSize()) threadsPerBlock <<= 1; // < 512 threads
+      
+      crossEntropyLossKernelOneBlock<<<1, threadsPerBlock, y.getSize() * sizeof(ftype)>>>(res.getData(), y.getData(), yPred.getData(), y.getSize());
+      cudaErrchk(cudaDeviceSynchronize());
+    }
+    else {
+      constexpr int threadsPerBlock = 256;
+      const int blocks = (y.getSize() + threadsPerBlock - 1) / threadsPerBlock;
+
+      ftype* tmp;
+      cudaErrchk(cudaMalloc(&tmp, blocks * sizeof(ftype)));
+
+      crossEntropyLossKernelOneBlock<<<blocks, threadsPerBlock, threadsPerBlock * sizeof(ftype)>>>(tmp, y.getData(), yPred.getData(), y.getSize());
+      cudaErrchk(cudaDeviceSynchronize());
+
+      // do a sum over the residual array
+      thrust::device_ptr<ftype> tmpPtr(tmp);
+      thrust::device_ptr<ftype> resPtr(res.getData());
+      resPtr[0] = thrust::reduce(tmpPtr, tmpPtr + blocks, static_cast<ftype>(0.0f), thrust::plus<ftype>());
+
+      cudaErrchk(cudaFree(tmp));
+    }
+
+    // loss = -loss / nBatches
+    const tensorSize_t stride = y.getDims()[-1];
+    const tensorSize_t nSamples = y.getSize() / stride;
+    divideScalarKernel<<<1, 1>>>(res.getData(), -1 * static_cast<ftype>(nSamples));
+    cudaErrchk(cudaDeviceSynchronize());
+  }
+
+  void crossEntropySoftmaxLoss(Tensor& res, const Tensor& y, const Tensor& yPred) {
+    const tensorSize_t stride   = static_cast<tensorSize_t>(yPred.getDims().get(-1));
+    const tensorSize_t nSamples = yPred.getSize() / stride;
+
+    ftype* maxValues;
+    cudaErrchk(cudaMalloc(&maxValues, nSamples * sizeof(ftype)));
+
+    // Find per-sample max values for numerical stability (mirrors softmax dispatch)
+    static const auto warpSizeT2 = 2 * DeviceProperties::getWarpSize();
+    if(stride <= warpSizeT2) {
+      constexpr int threadsPerBlock = 256;
+      const int blocks = (yPred.getSize() + threadsPerBlock - 1) / threadsPerBlock;
+
+      if(stride == 2)
+        findMaxKernelOneWarp<1><<<blocks, threadsPerBlock, threadsPerBlock * sizeof(ftype)>>>(maxValues, yPred.getData(), stride, yPred.getSize());
+      else if(stride <= 4)
+        findMaxKernelOneWarp<2><<<blocks, threadsPerBlock, threadsPerBlock * sizeof(ftype)>>>(maxValues, yPred.getData(), stride, yPred.getSize());
+      else if(stride <= 8)
+        findMaxKernelOneWarp<4><<<blocks, threadsPerBlock, threadsPerBlock * sizeof(ftype)>>>(maxValues, yPred.getData(), stride, yPred.getSize());
+      else if(stride <= 16)
+        findMaxKernelOneWarp<8><<<blocks, threadsPerBlock, threadsPerBlock * sizeof(ftype)>>>(maxValues, yPred.getData(), stride, yPred.getSize());
+      else if(stride <= 32)
+        findMaxKernelOneWarp<16><<<blocks, threadsPerBlock, threadsPerBlock * sizeof(ftype)>>>(maxValues, yPred.getData(), stride, yPred.getSize());
+      else
+        findMaxKernelOneWarp<32><<<blocks, threadsPerBlock, threadsPerBlock * sizeof(ftype)>>>(maxValues, yPred.getData(), stride, yPred.getSize());
+      cudaErrchk(cudaDeviceSynchronize());
+    }
+    else if(stride <= 512) {
+      int threadsPerBlock = 1;
+      while(threadsPerBlock < stride) threadsPerBlock <<= 1;
+      threadsPerBlock /= 2;
+
+      findMaxKernelOneBlock<<<nSamples, threadsPerBlock, 2 * threadsPerBlock * sizeof(ftype)>>>(maxValues, yPred.getData(), stride);
+      cudaErrchk(cudaDeviceSynchronize());
+    }
+    else {
+      cudaErrchk(cudaFree(maxValues));
+      __throw_invalid_argument("crossEntropySoftmaxLoss: stride > 512 not yet supported on CUDA");
+    }
+
+    // one block per sample, each thread covers up to 2 elements
+    int threadsPerBlock = 1;
+    while(threadsPerBlock * 2 < stride) threadsPerBlock <<= 1;
+    threadsPerBlock = max(1, threadsPerBlock);
+
+    ftype* perSampleLoss;
+    cudaErrchk(cudaMalloc(&perSampleLoss, nSamples * sizeof(ftype)));
+
+    crossEntropySoftmaxLossKernel<ftype><<<nSamples, threadsPerBlock, 2 * threadsPerBlock * sizeof(ftype)>>>(
+      perSampleLoss, y.getData(), yPred.getData(), maxValues, stride);
+    cudaErrchk(cudaDeviceSynchronize());
+
+    thrust::device_ptr<ftype> lossPtr(perSampleLoss);
+    thrust::device_ptr<ftype> resPtr(res.getData());
+    resPtr[0] = thrust::reduce(lossPtr, lossPtr + nSamples, static_cast<ftype>(0), thrust::plus<ftype>())
+                / static_cast<ftype>(nSamples);
+
+    cudaErrchk(cudaFree(perSampleLoss));
+    cudaErrchk(cudaFree(maxValues));
+  }
+
+  void rmseLoss(Tensor& res, const Tensor& y, const Tensor& yPred) {    
+    const auto nSamples = y.getSize();
+
+    if(nSamples <= 256) {
+      int threadsPerBlock = 1;
+      while(threadsPerBlock < nSamples) threadsPerBlock <<= 1; // < 512 threads
+      
+      rmseKernelOneBlock<<<1, threadsPerBlock, threadsPerBlock * sizeof(ftype)>>>(res.getData(), y.getData(), yPred.getData(), y.getSize());
+      cudaErrchk(cudaDeviceSynchronize());
+    }
+    else {
+      constexpr int threadsPerBlock = 256;
+      const int blocks = (nSamples + threadsPerBlock - 1) / threadsPerBlock;
+
+      ftype* tmp;
+      cudaErrchk(cudaMalloc(&tmp, blocks * sizeof(ftype)));
+
+      rmseKernelOneBlock<<<blocks, threadsPerBlock, threadsPerBlock * sizeof(ftype)>>>(tmp, y.getData(), yPred.getData(), y.getSize());
+      cudaErrchk(cudaDeviceSynchronize());
+
+      // do a sum over the residual array
+      thrust::device_ptr<ftype> tmpPtr(tmp);
+      thrust::device_ptr<ftype> resPtr(res.getData());
+      resPtr[0] = thrust::reduce(tmpPtr, tmpPtr + blocks, static_cast<ftype>(0.0f), thrust::plus<ftype>());
+
+      cudaErrchk(cudaFree(tmp));
+    }
+
+    normalizeRmse<ftype><<<1, 1>>>(res.getData(), static_cast<ftype>(nSamples));
+    cudaErrchk(cudaDeviceSynchronize());
+  }
+}
diff --git a/src/backend/training/loss_functions/cuda/loss_functions.cuh b/src/backend/training/loss_functions/cuda/loss_functions.cuh
new file mode 100644
index 0000000..fa41596
--- /dev/null
+++ b/src/backend/training/loss_functions/cuda/loss_functions.cuh
@@ -0,0 +1,28 @@
+/**
+ * @file loss_functions.cuh
+ * @author Robert Baumgartner (r.baumgartner-1@tudelft.nl)
+ * @brief
+ * @version 0.1
+ * @date 2026-05-10
+ *
+ * @copyright Copyright (c) 2026
+ *
+ */
+
+#pragma once
+
+#ifndef __CUDA
+static_assert(false, "File should not be included without CUDA enabled");
+#endif // __CUDA
+
+#include "data_modeling/tensor.h"
+
+namespace cuda_impl {
+  void bceLoss(Tensor& res, const Tensor& y, const Tensor& yPred);
+  void bceSigmoidLoss(Tensor& res, const Tensor& y, const Tensor& logits);
+
+  void crossEntropyLoss(Tensor& res, const Tensor& y, const Tensor& yPred);
+  void crossEntropySoftmaxLoss(Tensor& res, const Tensor& y, const Tensor& yPred);
+
+  void rmseLoss(Tensor& res, const Tensor& y, const Tensor& yPred);
+}
diff --git a/src/backend/training/loss_functions/rmse_loss.cpp b/src/backend/training/loss_functions/rmse_loss.cpp
index 30c750a..2684443 100644
--- a/src/backend/training/loss_functions/rmse_loss.cpp
+++ b/src/backend/training/loss_functions/rmse_loss.cpp
@@ -1,21 +1,26 @@
 /**
  * @file rmse_loss.cpp
  * @author Robert Baumgartner (r.baumgartner-1@tudelft.nl)
- * @brief 
+ * @brief
  * @version 0.1
  * @date 2026-03-14
- * 
+ *
  * @copyright Copyright (c) 2026
- * 
+ *
  */
 
 #include "rmse_loss.h"
-
 #include "computational_graph/loss_functions/rmse_node.h"
 
 #include <cmath>
 #include <iostream>
 
+#ifdef __CUDA
+#include "training/loss_functions/cuda/loss_functions.cuh"
+#else
+#include <stdexcept>
+#endif
+
 using namespace std;
 using namespace train;
 
@@ -26,7 +31,7 @@ using namespace train;
 shared_ptr<Tensor> RmseLoss::operator()(const shared_ptr<Tensor> y, const shared_ptr<Tensor> ypred) const {
   if(!ypred->getRequiresGrad()) {
     __throw_invalid_argument("ypred must have gradient enabled");
-  }  
+  }
   else if(y->getDevice() != ypred->getDevice()){
     __throw_invalid_argument("y and ypred must be on same device");
   }
@@ -34,22 +39,37 @@ shared_ptr<Tensor> RmseLoss::operator()(const shared_ptr<Tensor> y, const shared
     __throw_invalid_argument("Tensors must be of same shape");
   }
 
-  auto diffPow = [](ftype y, ftype ypred){
-    auto diff = y - ypred;
-    return diff * diff;
-  };
+  shared_ptr<Tensor> res = nullptr;
 
-  const auto nBatches = y->getDims()[0];
+  switch(y->getDevice()) {
+    case Device::CPU: {
+      auto diffPow = [](ftype y, ftype ypred){
+        auto diff = y - ypred;
+        return diff * diff;
+      };
 
-  ftype loss = 0;
-  for(tensorSize_t i=0; i<nBatches; i++){
-    loss += diffPow((*y)[i], (*ypred)[i]);
+      const auto nSamples = y->getSize();
+      
+      ftype loss = 0;
+      for(tensorSize_t i = 0; i < nSamples; i++){
+        loss += diffPow((*y)[i], (*ypred)[i]);
+      }
+      loss = sqrt(loss / nSamples);
+
+      res = make_shared<Tensor>(std::vector<tensorDim_t>{1}, std::vector<ftype>{loss}, y->getDevice(), true);
+      break;
+    }
+    case Device::CUDA:
+    #ifdef __CUDA
+      res = make_shared<Tensor>(vector<tensorDim_t>{1}, Device::CUDA, true);
+      cuda_impl::rmseLoss(*res, *y, *ypred);
+    #else
+      __throw_invalid_argument("Attempted to give CUDA tensor");
+    #endif
+      break;
   }
-  loss = sqrt(loss/nBatches);
 
-  auto res = make_shared<Tensor>(std::vector<tensorDim_t>{1}, std::vector<ftype>{loss}, y->getDevice(), true);
-  res->setCgNode(make_shared<cgraph::RmseNode>(y, ypred, loss));
+  res->setCgNode(make_shared<cgraph::RmseNode>(y, ypred, res->get(0)));
   assert(res->getRequiresGrad());
-
-  return res; 
-}
\ No newline at end of file
+  return res;
+}
diff --git a/src/backend/training/optimizers/cuda/optimizers.cu b/src/backend/training/optimizers/cuda/optimizers.cu
new file mode 100644
index 0000000..5038e93
--- /dev/null
+++ b/src/backend/training/optimizers/cuda/optimizers.cu
@@ -0,0 +1,67 @@
+/**
+ * @file optimizers.cu
+ * @author Robert Baumgartner (r.baumgartner-1@tudelft.nl)
+ * @brief
+ * @version 0.1
+ * @date 2026-05-13
+ *
+ * @copyright Copyright (c) 2026
+ *
+ */
+
+#ifndef __CUDA
+static_assert(false, "File should not be compiled without CUDA enabled");
+#endif // __CUDA
+
+#include "optimizers.cuh"
+#include "utility/cuda/cuda_common.cuh"
+#include "shared/cuda/common_kernels.cuh"
+
+#include "utility/global_params.h"
+
+using namespace std;
+
+namespace {
+  using namespace cuda_impl;
+  
+  __global__ void stepSgdKernel(ftype* const params, const ftype* const grads, const ftype lr, const tensorSize_t size) {
+    const int gid = blockIdx.x * blockDim.x + threadIdx.x;
+    if(gid >= size) {
+      return;
+    }
+
+    params[gid] = params[gid] - lr * grads[gid];
+  }
+
+  __global__ void stepRmsPropKernel(ftype* const tensor, ftype* const v, const ftype* const grads, const ftype lr, const ftype decay, const tensorSize_t size) {
+    const int gid = blockIdx.x * blockDim.x + threadIdx.x;
+    if(gid >= size) {
+      return;
+    }
+
+    const ftype g = grads[gid];
+    const ftype mavg = decay * v[gid] + (1 - decay) * g * g;
+    v[gid] = mavg;
+
+    const ftype update = tensor[gid] - (lr * g / (cudaSqrt<ftype>(mavg) + EPS_RMSPROP));
+    tensor[gid] = update;
+  }
+}
+
+namespace cuda_impl {
+  void sgdStep(Tensor& param, const Tensor& grads, ftype lr) {
+    constexpr int threadsPerBlock = 256;
+    const int blocks = (param.getSize() + threadsPerBlock - 1) / threadsPerBlock;
+
+    stepSgdKernel<<<blocks, threadsPerBlock>>>(param.getData(), grads.getData(), lr, param.getSize());
+    cudaErrchk(cudaDeviceSynchronize());
+  }
+
+  void rmspropStep(Tensor& param, Tensor& movingAvg, const Tensor& grads, ftype lr, ftype decay) {
+    constexpr int threadsPerBlock = 256;
+    const int blocks = (param.getSize() + threadsPerBlock - 1) / threadsPerBlock;
+
+    stepRmsPropKernel<<<blocks, threadsPerBlock>>>(param.getData(), movingAvg.getData(), grads.getData(), lr, decay, param.getSize());
+    cudaErrchk(cudaDeviceSynchronize());
+  }
+}
diff --git a/src/backend/training/optimizers/cuda/optimizers.cuh b/src/backend/training/optimizers/cuda/optimizers.cuh
new file mode 100644
index 0000000..a32562a
--- /dev/null
+++ b/src/backend/training/optimizers/cuda/optimizers.cuh
@@ -0,0 +1,24 @@
+/**
+ * @file optimizers.cuh
+ * @author Robert Baumgartner (r.baumgartner-1@tudelft.nl)
+ * @brief
+ * @version 0.1
+ * @date 2026-05-13
+ *
+ * @copyright Copyright (c) 2026
+ *
+ */
+
+#pragma once
+
+#ifndef __CUDA
+static_assert(false, "File should not be included without CUDA enabled");
+#endif // __CUDA
+
+#include "utility/global_params.h"
+#include "data_modeling/tensor.h"
+
+namespace cuda_impl {
+  void sgdStep(Tensor& param, const Tensor& grad, ftype lr);
+  void rmspropStep(Tensor& param, Tensor& movingAvg, const Tensor& grad, ftype lr, ftype decay);
+}
diff --git a/src/backend/training/optimizers/rmsprop.cpp b/src/backend/training/optimizers/rmsprop.cpp
index 6878f56..2d3f792 100644
--- a/src/backend/training/optimizers/rmsprop.cpp
+++ b/src/backend/training/optimizers/rmsprop.cpp
@@ -1,47 +1,70 @@
 /**
  * @file rmsprop.cpp
  * @author Robert Baumgartner (r.baumgartner-1@tudelft.nl)
- * @brief 
+ * @brief
  * @version 0.1
  * @date 2026-03-10
- * 
+ *
  * @copyright Copyright (c) 2026
- * 
+ *
  */
 
 #include "rmsprop.h"
 
+#ifdef __CUDA
+#include "training/optimizers/cuda/optimizers.cuh"
+#else
+#include <stdexcept>
+#endif
+
 using namespace std;
 using namespace train;
 
 void RmsPropOptimizer::step() {
-  constexpr ftype eps = 1e-8;
   for(const auto& param: params){
-    auto tPtr = param.get();
-    const auto gPtr = tPtr->getGrads().get();
-    auto vPtr = movingAvg[tPtr].get();
+    switch(param->getDevice()) {
+      case Device::CPU: {
+        auto tPtr = param.get();
+        const auto gPtr = tPtr->getGrads().get();
+
+        auto vPtr = movingAvg[tPtr].get();
+
+        if(vPtr != nullptr) [[likely]] {
+          for(tensorSize_t i = 0; i < gPtr->getSize(); i++){
+            auto g = (*gPtr)[i];
+            auto update = decay * (*vPtr)[i] + (1 - decay) * g * g;
+            vPtr->set(update, i);
+          }
+        }
+        else [[unlikely]] {
+          movingAvg[tPtr] = make_unique<Tensor>(tPtr->getDims(), Device::CPU, false);
+          vPtr = movingAvg[tPtr].get();
+          for(tensorSize_t i = 0; i < tPtr->getSize(); i++) {
+            auto g = (*gPtr)[i];
+            vPtr->set((1 - decay) * g * g, i);
+          }
+        }
 
-    // update moving avg
-    if(vPtr!=nullptr) { // hot path
-      for(tensorSize_t i=0; i<gPtr->getSize(); i++){ 
-        auto g = (*gPtr)[i];
-        auto update = decay * (*vPtr)[i] + (1-decay)*g*g;
-        vPtr->set(update, i);
+        for(tensorSize_t i = 0; i < tPtr->getSize(); i++) {
+          auto update = (*tPtr)[i] - (lr * (*gPtr)[i] / (sqrt((*vPtr)[i]) + EPS_RMSPROP));
+          tPtr->set(update, i);
+        }
+        break;
       }
-    }
-    else { // init loop
-      movingAvg[tPtr] = make_unique<Tensor>(tPtr->getDims(), tPtr->getDevice(), false); // create empty tensor
-      vPtr = movingAvg[tPtr].get();
-      for(tensorSize_t i=0; i<tPtr->getSize(); i++) {
-        auto g = (*gPtr)[i];
-        vPtr->set((1-decay)*g*g, i);
+      case Device::CUDA:
+      #ifdef __CUDA
+      {
+        auto tPtr = param.get();
+        if(movingAvg[tPtr] == nullptr) {
+          movingAvg[tPtr] = make_unique<Tensor>(tPtr->getDims(), Device::CUDA, false);
+          movingAvg[tPtr]->reset(0);
+        }
+        cuda_impl::rmspropStep(*param, *movingAvg[tPtr], *(param->getGrads()), lr, decay);
       }
-    }
-
-    // update gradients
-    for(tensorSize_t i=0; i<tPtr->getSize(); i++) {
-      auto update = (*tPtr)[i] - lr * (*gPtr)[i] / ((*vPtr)[i] + eps);
-      tPtr->set(update, i);
+      #else
+        __throw_invalid_argument("Attempted to use CUDA tensor");
+      #endif
+        break;
     }
   }
-}
\ No newline at end of file
+}
diff --git a/src/backend/training/optimizers/sgd.cpp b/src/backend/training/optimizers/sgd.cpp
index 83f122d..c3dc0f6 100644
--- a/src/backend/training/optimizers/sgd.cpp
+++ b/src/backend/training/optimizers/sgd.cpp
@@ -1,25 +1,42 @@
 /**
  * @file sgd.cpp
  * @author Robert Baumgartner (r.baumgartner-1@tudelft.nl)
- * @brief 
+ * @brief
  * @version 0.1
  * @date 2026-03-08
- * 
+ *
  * @copyright Copyright (c) 2026
- * 
+ *
  */
 
 #include "sgd.h"
 
+#ifdef __CUDA
+#include "training/optimizers/cuda/optimizers.cuh"
+#else
+#include <stdexcept>
+#endif
+
 using namespace std;
 using namespace train;
 
 void SgdOptimizer::step() {
   for(auto& t: params){
     auto grads = t->getGrads();
-    for(auto idx=0; idx<t->getSize(); idx++){
-      auto updatedWeight = (*t)[idx] - lr*(*grads)[idx];
-      t->set(updatedWeight, idx);
+    switch(t->getDevice()) {
+      case Device::CPU:
+        for(auto idx = 0; idx < t->getSize(); idx++){
+          auto updatedWeight = (*t)[idx] - lr * (*grads)[idx];
+          t->set(updatedWeight, idx);
+        }
+        break;
+      case Device::CUDA:
+      #ifdef __CUDA
+        cuda_impl::sgdStep(*t, *grads, lr);
+      #else
+        __throw_invalid_argument("Attempted to use CUDA tensor");
+      #endif
+        break;
     }
   }
-}
\ No newline at end of file
+}
diff --git a/src/backend/utility/cuda/cuda_common.cu b/src/backend/utility/cuda/cuda_common.cu
new file mode 100644
index 0000000..101ea3f
--- /dev/null
+++ b/src/backend/utility/cuda/cuda_common.cu
@@ -0,0 +1,28 @@
+/**
+ * @file cuda_common.cu
+ * @author Robert Baumgartner (r.baumgartner-1@tudelft.nl)
+ * @brief 
+ * @version 0.1
+ * @date 2026-03-22
+ * 
+ * @copyright Copyright (c) 2026
+ * 
+ */
+
+#ifndef __CUDA
+static_assert(false, "File should not be included without CUDA enabled");
+#endif
+
+#include "cuda_common.cuh"
+
+#include <iostream>
+
+void utility::gpuAssert(cudaError_t code, const char *file, int line, bool abort)
+{
+  if (code != cudaSuccess) 
+  {
+      //fprintf(stderr,"GPUassert: %s %s %d\n", cudaGetErrorString(code), file, line);
+      std::cerr << "GPUassert: " << cudaGetErrorString(code) << " " << file << " " << line;
+      if (abort) exit(code);
+  }
+}
\ No newline at end of file
diff --git a/src/backend/utility/cuda/cuda_common.cuh b/src/backend/utility/cuda/cuda_common.cuh
new file mode 100644
index 0000000..7d505d5
--- /dev/null
+++ b/src/backend/utility/cuda/cuda_common.cuh
@@ -0,0 +1,60 @@
+/**
+ * @file cuda_common.cuh
+ * @author Robert Baumgartner (r.baumgartner-1@tudelft.nl)
+ * @brief 
+ * @version 0.1
+ * @date 2026-03-22
+ * 
+ * @copyright Copyright (c) 2026
+ * 
+ */
+
+#pragma once
+
+#ifndef __CUDA
+static_assert(false, "File should not be included without CUDA enabled");
+#endif // __CUDA
+
+#include "cuda_runtime.h"
+#include "utility/global_params.h"
+
+#include <type_traits>
+
+namespace utility {
+  void gpuAssert(cudaError_t code, const char *file, int line, bool abort=true);
+}
+
+#define cudaErrchk(ans) { utility::gpuAssert((ans), __FILE__, __LINE__); }
+
+namespace cuda_impl {
+  struct DeviceProperties final {
+    private:
+      int threadsPerBlock;
+      int warpSize;
+
+      static const DeviceProperties& get() {
+        static DeviceProperties instance;
+        return instance;
+      }
+
+      DeviceProperties() {
+        cudaDeviceProp prop;
+        cudaGetDeviceProperties(&prop, 0);
+
+        threadsPerBlock = prop.maxThreadsPerBlock;
+        warpSize = prop.warpSize;
+      }
+
+    public:
+      DeviceProperties(const DeviceProperties&) = delete;
+      DeviceProperties& operator=(const DeviceProperties&) = delete;
+      
+      DeviceProperties(DeviceProperties&&) = delete;
+      DeviceProperties& operator=(DeviceProperties&&) = delete;
+      
+      ~DeviceProperties() = default;
+
+      static int getThreadsPerBlock() { return get().threadsPerBlock; }
+      static int getWarpSize() { return get().warpSize; }
+  };
+}
diff --git a/src/backend/utility/global_params.h b/src/backend/utility/global_params.h
index 3d6edcb..b2861e5 100644
--- a/src/backend/utility/global_params.h
+++ b/src/backend/utility/global_params.h
@@ -13,6 +13,8 @@
 
 #include <cstdint>
 
+#include <array>
+
 using ftype = float; // TODO: make compiler flag?
 
 /** 
@@ -26,14 +28,36 @@ using ftype = float; // TODO: make compiler flag?
  * FOR OVERFLOWS. Similarly, we assume that all dimensions you
  * request fit into datatype tensorDim_t.
  */ 
-using tensorDim_t = std::uint16_t;
+using tensorDim_t = std::uint32_t;
 using tensorSize_t = std::uint32_t;
 
+/**
+ * Note: If we want to narrow down tensorDim_t, say for embedded
+ * devices, we can easily run into subtle bugs, as overflows of 
+ * integers are not detected. A quick way out could be the 
+ * following struct, which is currently not implemented yet.
+ * 
+ * struct tensorDim_t {
+ *   uint16_t value;
+ *   tensorSize_t operator*(const tensorDim_t& other) const {
+ *     return static_cast<tensorSize_t>(value) * other.value;
+ *   }
+ * };
+ */
+
+constexpr tensorDim_t MAX_NDIMS = 6;
+
 // we assert this here so during conversions of tensorDim_t to 
 // tensorSize_t we do not need to cast explicitely
 static_assert(sizeof(tensorDim_t)<=sizeof(tensorSize_t));
 
 // ----------------- Numerical stability -------------------
 
-constexpr ftype epsCrossentropy = 1e-5;
-constexpr ftype epsBce = 1e-5;
\ No newline at end of file
+constexpr ftype EPS_CROSSENTROPY = 1e-5;
+constexpr ftype EPS_BCE = 1e-5;
+constexpr ftype EPS_RMSE = 1e-9;
+constexpr ftype EPS_RMSPROP = 1e-8;
+
+// ----------------- Default values ------------------------
+
+constexpr ftype EPS_LEAKY_RELU = 0.01;
diff --git a/src/backend/utility/utils.h b/src/backend/utility/utils.h
new file mode 100644
index 0000000..660a4bf
--- /dev/null
+++ b/src/backend/utility/utils.h
@@ -0,0 +1,39 @@
+/**
+ * @file utils.h
+ * @author Robert Baumgartner (r.baumgartner-1@tudelft.nl)
+ * @brief 
+ * @version 0.1
+ * @date 2026-05-11
+ * 
+ * @copyright Copyright (c) 2026
+ * 
+ */
+
+#pragma once 
+
+#ifdef NDEBUG
+  #define assert_debug(cond, msg) ((void)0)
+#else
+  #define assert_debug(cond, msg) \
+    do { \
+      if (!(cond)) { \
+        fprintf(stderr, "Assertion failed: %s\n  Message: %s\n  File: %s, Line: %d\n", \
+                #cond, msg, __FILE__, __LINE__); \
+        abort(); \
+      } \
+    } while(0)
+#endif
+
+template <class... T>
+constexpr bool always_false = false;
+
+template<typename T>
+struct FtypeWarning {
+    static constexpr void check() {}
+};
+
+template<>
+struct FtypeWarning<double> {
+    [[deprecated("ftype=double has serious CUDA performance implications")]]
+    static constexpr void check() {}
+};
\ No newline at end of file
diff --git a/src/python/py_core/py_core.cpp b/src/python/py_core/py_core.cpp
index 9c7fab2..8c1bec8 100644
--- a/src/python/py_core/py_core.cpp
+++ b/src/python/py_core/py_core.cpp
@@ -19,6 +19,8 @@
 #include "data_modeling/tensor_functions.h"
 #include "computational_graph/tensor_ops/graph_creation.h"
 
+#include "utility/utils.h"
+
 #include <boost/python.hpp>
 #include <boost/python/enum.hpp>
 #include <boost/python/return_internal_reference.hpp>
@@ -28,6 +30,8 @@
 
 BOOST_PYTHON_MODULE(_core)
 {
+  FtypeWarning<ftype>::check();
+
   using namespace boost::python;
 
   // some macros to make code below easier to read
@@ -47,8 +51,13 @@ BOOST_PYTHON_MODULE(_core)
   }
 
   // different, since those are not methods anymore
+  #define WRAP_FREE_MEMBER_FUNC_0(fPtr) \
+  +[](const Tensor& self) -> std::shared_ptr<Tensor> { \
+    return std::make_shared<Tensor>((self.*fPtr)()); \
+  }
+
   #define WRAP_FREE_MEMBER_FUNC_1(fPtr, T1, T2) \
-  +[](const Tensor& self, int v1, int v2) -> std::shared_ptr<Tensor> { \
+  +[](const Tensor& self, T1 v1, T2 v2) -> std::shared_ptr<Tensor> { \
     return std::make_shared<Tensor>((self.*fPtr)(v1, v2)); \
   }
 
@@ -155,21 +164,21 @@ BOOST_PYTHON_MODULE(_core)
         
     // static creation methods
     .def("ones", WRAP_FREE_FUNC_1(Py_DataModeling::Ones0, std::vector<tensorDim_t>))
-    .def("ones", WRAP_FREE_FUNC_2(Py_DataModeling::Ones1, std::vector<tensorDim_t>, Device))
     .def("ones", WRAP_FREE_FUNC_2(Py_DataModeling::Ones2, std::vector<tensorDim_t>, const bool))
+    .def("ones", WRAP_FREE_FUNC_2(Py_DataModeling::Ones1, std::vector<tensorDim_t>, Device))
     .def("ones", WRAP_FREE_FUNC_3(Py_DataModeling::Ones3, std::vector<tensorDim_t>, Device, const bool))
     .staticmethod("ones")
 
     .def("zeros", WRAP_FREE_FUNC_1(Py_DataModeling::Zeros0, std::vector<tensorDim_t>))
-    .def("zeros", WRAP_FREE_FUNC_2(Py_DataModeling::Zeros1, std::vector<tensorDim_t>, Device))
     .def("zeros", WRAP_FREE_FUNC_2(Py_DataModeling::Zeros2, std::vector<tensorDim_t>, const bool))
+    .def("zeros", WRAP_FREE_FUNC_2(Py_DataModeling::Zeros1, std::vector<tensorDim_t>, Device))
     .def("zeros", WRAP_FREE_FUNC_3(Py_DataModeling::Zeros3, std::vector<tensorDim_t>, Device, const bool))
     .staticmethod("zeros")
 
     .def("gauss", WRAP_FREE_FUNC_2(Py_DataModeling::Gaussian0, std::vector<tensorDim_t>, ftype))
-    .def("gauss", WRAP_FREE_FUNC_3(Py_DataModeling::Gaussian1, std::vector<tensorDim_t>, Device, ftype))
     .def("gauss", WRAP_FREE_FUNC_3(Py_DataModeling::Gaussian2, std::vector<tensorDim_t>, ftype, const bool))
-    .def("gauss", WRAP_FREE_FUNC_8(Py_DataModeling::Gaussian3, std::vector<tensorDim_t>, Device, ftype, const bool))
+    .def("gauss", WRAP_FREE_FUNC_3(Py_DataModeling::Gaussian1, std::vector<tensorDim_t>, ftype, Device))
+    .def("gauss", WRAP_FREE_FUNC_8(Py_DataModeling::Gaussian3, std::vector<tensorDim_t>, ftype, Device, const bool))
     .staticmethod("gauss")
 
     // properties
@@ -217,10 +226,12 @@ BOOST_PYTHON_MODULE(_core)
       return t->hasGrads();
       })
 
-    .def("transpose", WRAP_FREE_MEMBER_FUNC_1(Py_DataModeling::transpose1, int, int))
-    .def("transpose", WRAP_FREE_MEMBER_FUNC_2(Py_DataModeling::transpose2, int, int, bool))
-    .def("transposeThis", Py_DataModeling::transposeThis1)
-    .def("transposeThis", Py_DataModeling::transposeThis2)
+    .def("transpose", +[](Tensor& self) -> std::shared_ptr<Tensor> {
+        return std::make_shared<Tensor>(self.transpose());
+      })
+    .def("transpose", +[](Tensor& self, int dim1, int dim2) -> std::shared_ptr<Tensor> {
+        return std::make_shared<Tensor>(self.transpose(dim1, dim2));
+      })
         
     .def("backward", &Tensor::backward)
   ;
@@ -237,9 +248,9 @@ BOOST_PYTHON_MODULE(_core)
   def("Zeros", WRAP_FREE_FUNC_3(Py_DataModeling::Zeros3, std::vector<tensorDim_t>, Device, const bool));
 
   def("Gaussian", WRAP_FREE_FUNC_2(Py_DataModeling::Gaussian0, std::vector<tensorDim_t>, ftype));
-  def("Gaussian", WRAP_FREE_FUNC_3(Py_DataModeling::Gaussian1, std::vector<tensorDim_t>, Device, ftype));
+  def("Gaussian", WRAP_FREE_FUNC_3(Py_DataModeling::Gaussian1, std::vector<tensorDim_t>, ftype, Device));
   def("Gaussian", WRAP_FREE_FUNC_3(Py_DataModeling::Gaussian2, std::vector<tensorDim_t>, ftype, const bool));
-  def("Gaussian", WRAP_FREE_FUNC_8(Py_DataModeling::Gaussian3, std::vector<tensorDim_t>, Device, ftype, const bool));
+  def("Gaussian", WRAP_FREE_FUNC_8(Py_DataModeling::Gaussian3, std::vector<tensorDim_t>, ftype, Device, const bool));
 
   // must call this before using numpy C API
   Py_DataModeling::initNumpy();  // numpy initialization
diff --git a/src/python/py_core/py_core_util.h b/src/python/py_core/py_core_util.h
index 4da52cf..e6c7b21 100644
--- a/src/python/py_core/py_core_util.h
+++ b/src/python/py_core/py_core_util.h
@@ -78,8 +78,8 @@ namespace Py_DataModeling
     return TensorFunctions::Gaussian(std::move(dims), stddev);
   }
 
-  inline auto GaussianWrapper1(std::vector<tensorDim_t> dims, Device d, ftype stddev) { 
-    return TensorFunctions::Gaussian(std::move(dims), d, stddev);
+  inline auto GaussianWrapper1(std::vector<tensorDim_t> dims, ftype stddev, Device d) { 
+    return TensorFunctions::Gaussian(std::move(dims), stddev, d);
   }
 
   inline Tensor    (*Ones0)(std::vector<tensorDim_t>)                                             = &OnesWrapper0;
@@ -93,19 +93,16 @@ namespace Py_DataModeling
   inline Tensor    (*Zeros3)(std::vector<tensorDim_t>, Device, const bool)                        = &(TensorFunctions::Zeros);
 
   inline Tensor    (*Gaussian0)(std::vector<tensorDim_t>, ftype)                                  = &GaussianWrapper0;
-  inline Tensor    (*Gaussian1)(std::vector<tensorDim_t>, Device, ftype)                          = &GaussianWrapper1;
+  inline Tensor    (*Gaussian1)(std::vector<tensorDim_t>, ftype, Device)                          = &GaussianWrapper1;
   inline Tensor    (*Gaussian2)(std::vector<tensorDim_t>, ftype, const bool)                      = &(TensorFunctions::Gaussian);
-  inline Tensor    (*Gaussian3)(std::vector<tensorDim_t>, Device, ftype, const bool)              = &(TensorFunctions::Gaussian);
+  inline Tensor    (*Gaussian3)(std::vector<tensorDim_t>, ftype, Device, const bool)              = &(TensorFunctions::Gaussian);
 
   inline void    (Tensor::*reset1)(const ftype)                                                   = &Tensor::reset;
   inline void    (Tensor::*reset2)(const std::shared_ptr<utility::InitializerBase>)               = &Tensor::reset;
 
-  inline void    (Tensor::*transposeThis1)()                                                      = &Tensor::transposeThis;
-  inline void    (Tensor::*transposeThis2)(int, int)                                              = &Tensor::transposeThis;
-  inline Tensor  (Tensor::*transpose1)(int, int) const                                            = &Tensor::transpose;
-  inline Tensor  (Tensor::*transpose2)(int, int, bool) const                                      = &Tensor::transpose;
+  inline Tensor  (Tensor::*transpose2)(int, int)                                                  = &Tensor::transpose;
 
-  inline ftype   (Tensor::*getItemVector)(const std::vector<tensorDim_t>&) const                       = &Tensor::get;
+  inline ftype   (Tensor::*getItemVector)(const std::vector<tensorDim_t>&) const                  = &Tensor::get;
 
   /*********************************************************************************************************
   ***************************************** Graph creation *************************************************
diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt
index c8ca76b..08fb29c 100644
--- a/tests/CMakeLists.txt
+++ b/tests/CMakeLists.txt
@@ -1,38 +1,77 @@
 include(FetchContent)
+
 FetchContent_Declare(
-    googletest
-    GIT_REPOSITORY https://github.com/google/googletest.git
-    GIT_TAG v1.14.0
+  googletest
+  GIT_REPOSITORY https://github.com/google/googletest.git
+  GIT_TAG v1.14.0
 )
 FetchContent_MakeAvailable(googletest)
 
 add_executable(unit_tests_backend
-    backend/test_data_modeling.cpp
-    backend/test_computational_graph.cpp
-    backend/test_module.cpp
-    backend/test_losses.cpp
-    backend/test_train_loop.cpp
-    )
+  backend/main.cpp
+  backend/test_tensorops.cpp
+  backend/test_module.cpp
+  backend/test_losses.cpp
+  backend/test_computational_graph.cpp
+  )
 
 target_link_libraries(unit_tests_backend PRIVATE
-    gtest_main # pre-built main, avoids boilerplate if no custom initialization needed
-    BackendCore
+  gtest  
+  BackendCore
 )
 
 include(GoogleTest)
 gtest_discover_tests(unit_tests_backend)
 
-find_package(Python3 COMPONENTS Interpreter)
-if(Python3_FOUND)
-    # replace the placeholder variables and copy resulting file in .py file
+if(CMAKE_CUDA_COMPILER)
+  file(GLOB_RECURSE CUDA_TEST_SOURCES backend/cuda/*.cu)
+  add_executable(unit_tests_cuda ${CUDA_TEST_SOURCES})
+
+  find_package(Threads REQUIRED)
+  target_link_libraries(unit_tests_cuda PRIVATE
+    gtest
+    Threads::Threads # platform agnostic
+    BackendCore
+    CUDA::cudart
+  )
+
+  set_target_properties(unit_tests_cuda PROPERTIES
+    CUDA_SEPARABLE_COMPILATION ON
+    CMAKE_CUDA_ARCHITECTURES native
+  )
+
+  find_package(CUDAToolkit REQUIRED)
+  target_include_directories(unit_tests_cuda PRIVATE
+    ${CUDAToolkit_INCLUDE_DIRS}
+    ${CMAKE_CURRENT_SOURCE_DIR}/backend
+  )
+
+  gtest_discover_tests(unit_tests_cuda)
+
+  find_package(Python3 COMPONENTS Interpreter)
+  if(Python3_FOUND)
     add_test(
-        NAME python_tests
-        COMMAND ${Python3_EXECUTABLE} -m pytest 
-                ${CMAKE_CURRENT_SOURCE_DIR}/python
+      NAME python_tests_cuda
+      COMMAND ${Python3_EXECUTABLE} -m pytest
+              ${CMAKE_CURRENT_SOURCE_DIR}/python/cuda
     )
-
-    # Set environment for Python to find the module
-    set_tests_properties(python_tests PROPERTIES
-        ENVIRONMENT "PYTHONPATH=${PYTHON_MODULE_DIR}:$ENV{PYTHONPATH}"
+    set_tests_properties(python_tests_cuda PROPERTIES
+      ENVIRONMENT "PYTHONPATH=${PYTHON_MODULE_DIR}:$ENV{PYTHONPATH}"
     )
+  endif()
+endif()
+
+find_package(Python3 COMPONENTS Interpreter)
+if(Python3_FOUND)
+  # replace the placeholder variables and copy resulting file in .py file
+  add_test(
+    NAME python_tests
+    COMMAND ${Python3_EXECUTABLE} -m pytest 
+            ${CMAKE_CURRENT_SOURCE_DIR}/python
+  )
+
+  # Set environment for Python to find the module
+  set_tests_properties(python_tests PROPERTIES
+    ENVIRONMENT "PYTHONPATH=${PYTHON_MODULE_DIR}:$ENV{PYTHONPATH}"
+  )
 endif()
\ No newline at end of file
diff --git a/tests/backend/cuda/main.cu b/tests/backend/cuda/main.cu
new file mode 100644
index 0000000..8cca788
--- /dev/null
+++ b/tests/backend/cuda/main.cu
@@ -0,0 +1,36 @@
+/**
+ * @file main.cu
+ * @author Robert Baumgartner (r.baumgartner-1@tudelft.nl)
+ * @brief 
+ * @version 0.1
+ * @date 2026-05-12
+ * 
+ * @copyright Copyright (c) 2026
+ * 
+ */
+
+#ifndef __CUDA
+static_assert(false, "File should not be compiled without CUDA enabled");
+#endif // __CUDA
+
+#include <gtest/gtest.h>
+
+#include "system/sys_functions.h"
+
+class CudaEnvironment : public ::testing::Environment {
+public:
+  void SetUp() override {
+    // cuda warmup to avoid context initialization costs
+    void* tmp;
+    cudaMalloc(&tmp, 1);
+    cudaFree(tmp);
+    cudaDeviceSynchronize();
+  }
+};
+
+int main(int argc, char** argv) {
+  ::testing::InitGoogleTest(&argc, argv);
+  ::testing::AddGlobalTestEnvironment(new CudaEnvironment());
+  sys::setRandomSeed(42);
+  return RUN_ALL_TESTS();
+}
\ No newline at end of file
diff --git a/tests/backend/cuda/netfacories.cuh b/tests/backend/cuda/netfacories.cuh
new file mode 100644
index 0000000..faba0cd
--- /dev/null
+++ b/tests/backend/cuda/netfacories.cuh
@@ -0,0 +1,2 @@
+#pragma once
+#include "net_factories.h"
diff --git a/tests/backend/cuda/test_computational_graph.cu b/tests/backend/cuda/test_computational_graph.cu
new file mode 100644
index 0000000..31950da
--- /dev/null
+++ b/tests/backend/cuda/test_computational_graph.cu
@@ -0,0 +1,247 @@
+/**
+ * @file test_computational_graph.cu
+ * @author Robert Baumgartner (r.baumgartner-1@tudelft.nl)
+ * @brief
+ * @version 0.1
+ * @date 2026-05-20
+ *
+ * @copyright Copyright (c) 2026
+ *
+ */
+
+#ifndef __CUDA
+static_assert(false, "File should not be compiled without CUDA enabled");
+#endif // __CUDA
+
+#include <gtest/gtest.h>
+
+#include "data_modeling/tensor_functions.h"
+#include "computational_graph/tensor_ops/graph_creation.h"
+
+#include "net_factories.h"
+
+#include "training/optimizers/sgd.h"
+#include "training/optimizers/rmsprop.h"
+
+#include "training/loss_functions/bce_loss.h"
+#include "training/loss_functions/crossentropy_loss.h"
+#include "training/loss_functions/bce_sigmoid_loss.h"
+#include "training/loss_functions/crossentropy_softmax_loss.h"
+
+#include "training/trainers/base_train_loop.h"
+
+using namespace std;
+
+TEST(CudaAutogradTest, ThrowsIfNoGradientSet) {
+  auto t1 = TensorFunctions::makeSharedTensor({1}, {3.0}, Device::CUDA, false);
+  auto t2 = TensorFunctions::makeSharedTensor({1}, {2.0}, Device::CUDA, false);
+
+  auto loss = cgraph::add(t1, t2);
+
+  ASSERT_THROW(loss->backward(), std::runtime_error);
+}
+
+TEST(CudaAutogradTest, ChainRule) {
+  auto x = TensorFunctions::makeSharedTensor({1}, {2.0}, Device::CUDA, true);
+
+  auto y = cgraph::mul(x, x);      // y = x^2
+  auto z = cgraph::add(x, y);      // z = x^2 + x
+  auto loss = cgraph::mul(z, z);   // loss = (x^2 + x)^2
+
+  loss->backward();
+
+  // dloss/dx = 2(x^2 + x) * (2x + 1)
+  // At x=2: 2(4 + 2) * (4 + 1) = 60
+  ASSERT_NEAR(x->getGrads()->get(0), 60.0, 1e-4);
+}
+
+TEST(CudaAutogradTest, MultiVariateChainRule) {
+  auto x = TensorFunctions::makeSharedTensor({2}, {1.0, 2.0}, Device::CUDA, true);
+
+  auto y = cgraph::mul(x, 3.0); // y = [3, 6]
+  auto loss = TensorFunctions::makeSharedTensor({1}, {0.0}, Device::CUDA, true);
+  for(int i = 0; i < y->getSize(); i++){
+    loss = cgraph::add(loss, cgraph::get(y, i));
+  }    // loss = 9
+
+  loss->backward();
+
+  // dloss/dx = 3
+  ASSERT_NEAR(x->getGrads()->get(0), 3.0, 1e-4);
+  ASSERT_NEAR(x->getGrads()->get(1), 3.0, 1e-4);
+
+  ASSERT_NEAR(y->getGrads()->get(0), 1.0, 1e-4);
+  ASSERT_NEAR(y->getGrads()->get(1), 1.0, 1e-4);
+}
+
+TEST(CudaOverfitTest, BceSgdOverfitsSmallDataset) {
+  auto x = TensorFunctions::makeSharedTensor(
+    {4, 2}, {0.0, 0.0,
+             0.0, 1.0,
+             1.0, 0.0,
+             1.0, 1.0}, Device::CUDA, false);
+
+  auto y = TensorFunctions::makeSharedTensor(
+    {4, 1}, {0.0,
+             1.0,
+             1.0,
+             0.0}, Device::CUDA, false);
+
+  auto net = makeBinaryNet(Device::CUDA);
+  auto loss = make_shared<train::BceLoss>();
+  auto optim = make_shared<train::SgdOptimizer>(
+    net->parameters(), /*lr=*/0.05);
+
+  auto trainLoop = train::BaseTrainLoop(
+    net, loss, optim, /*epochs=*/2000, /*bsize=*/static_cast<tensorDim_t>(4));
+
+  trainLoop.run(x, y, /*shuffle=*/false, /*verbose=*/false);
+
+  auto pred = (*net)(x);
+  auto finalLoss = (*loss)(y, pred);
+
+  EXPECT_LT((*finalLoss)[0], 0.05f)
+    << "Network failed to overfit binary dataset\n"
+    << "Final loss: " << *finalLoss;
+}
+
+TEST(CudaOverfitTest, BceSgdOverfitsSmallDataset_OptimizedLoss) {
+  auto x = TensorFunctions::makeSharedTensor(
+    {4, 2}, {0.0, 0.0,
+             0.0, 1.0,
+             1.0, 0.0,
+             1.0, 1.0}, Device::CUDA, false);
+
+  auto y = TensorFunctions::makeSharedTensor(
+    {4, 1}, {0.0,
+             1.0,
+             1.0,
+             0.0}, Device::CUDA, false);
+
+  auto net = makeBinaryNet2(Device::CUDA);
+  auto loss = make_shared<train::BceSigmoidLoss>();
+  auto optim = make_shared<train::SgdOptimizer>(
+    net->parameters(), /*lr=*/0.05);
+
+  auto trainLoop = train::BaseTrainLoop(
+    net, loss, optim, /*epochs=*/2000, /*bsize=*/static_cast<tensorDim_t>(4));
+
+  trainLoop.run(x, y, /*shuffle=*/false, /*verbose=*/false);
+
+  auto pred = (*net)(x);
+  auto finalLoss = (*loss)(y, pred);
+
+  EXPECT_LT((*finalLoss)[0], 0.05f)
+    << "Network failed to overfit binary dataset\n"
+    << "Final loss: " << *finalLoss;
+}
+
+TEST(CudaOverfitTest, CrossEntropyRMSPropOverfitsSmallDataset) {
+  auto x = TensorFunctions::makeSharedTensor(
+    {6, 2}, {1.0, 0.0,
+             1.0, 0.1,
+             0.0, 1.0,
+             0.1, 1.0,
+             0.5, 0.5,
+             0.4, 0.6}, Device::CUDA, false);
+
+  auto y = TensorFunctions::makeSharedTensor(
+    {6, 3}, {1.0, 0.0, 0.0,
+             1.0, 0.0, 0.0,
+             0.0, 1.0, 0.0,
+             0.0, 1.0, 0.0,
+             0.0, 0.0, 1.0,
+             0.0, 0.0, 1.0}, Device::CUDA, false);
+
+  auto net = makeMulticlassNet(Device::CUDA);
+  auto loss = make_shared<train::CrossEntropyLoss>();
+  auto optim = make_shared<train::RmsPropOptimizer>(
+    net->parameters(), /*lr=*/0.001, /*decay=*/0.95);
+
+  auto trainLoop = train::BaseTrainLoop(
+    net, loss, optim, /*epochs=*/2000, /*bsize=*/6);
+
+  trainLoop.run(x, y, /*shuffle=*/false, /*verbose=*/false);
+
+  auto pred = (*net)(x);
+  auto finalLoss = (*loss)(y, pred);
+
+
+  EXPECT_LT((*finalLoss)[0], 0.05f)
+    << "Network failed to overfit multiclass dataset\n"
+    << "Final loss: " << *finalLoss;
+}
+
+TEST(CudaOverfitTest, CrossEntropyRMSPropOverfitsSmallDataset_OptimizedLoss) {
+  auto x = TensorFunctions::makeSharedTensor(
+    {6, 2}, {1.0, 0.0,
+             1.0, 0.1,
+             0.0, 1.0,
+             0.1, 1.0,
+             0.5, 0.5,
+             0.4, 0.6}, Device::CUDA, false);
+
+  auto y = TensorFunctions::makeSharedTensor(
+    {6, 3}, {1.0, 0.0, 0.0,
+             1.0, 0.0, 0.0,
+             0.0, 1.0, 0.0,
+             0.0, 1.0, 0.0,
+             0.0, 0.0, 1.0,
+             0.0, 0.0, 1.0}, Device::CUDA, false);
+
+  auto net = makeMulticlassNet2(Device::CUDA);
+  auto loss = make_shared<train::CrossEntropySoftmaxLoss>();
+  auto optim = make_shared<train::RmsPropOptimizer>(
+    net->parameters(), /*lr=*/0.001, /*decay=*/0.95);
+
+  auto trainLoop = train::BaseTrainLoop(
+    net, loss, optim, /*epochs=*/2000, /*bsize=*/6);
+
+  trainLoop.run(x, y, /*shuffle=*/false, /*verbose=*/false);
+
+  auto pred = (*net)(x);
+  auto finalLoss = (*loss)(y, pred);
+
+  EXPECT_LT((*finalLoss)[0], 0.05f)
+    << "Network failed to overfit multiclass dataset\n"
+    << "Final loss: " << *finalLoss;
+}
+
+TEST(CudaOptimizerTest, ZeroGrad_ClearsAllGradients) {
+  auto x = TensorFunctions::makeSharedTensor(
+    {4, 2}, {0.0, 0.0, 0.0, 1.0, 1.0, 0.0, 1.0, 1.0}, Device::CUDA, false);
+  auto y = TensorFunctions::makeSharedTensor(
+    {4, 1}, {0.0, 1.0, 1.0, 0.0}, Device::CUDA, false);
+
+  auto net = makeBinaryNet(Device::CUDA);
+  auto loss = make_shared<train::BceSigmoidLoss>();
+  auto optim = make_shared<train::SgdOptimizer>(net->parameters(), 0.01f);
+
+  auto pred = (*net)(x);
+  auto l = (*loss)(y, pred);
+  l->backward();
+
+  bool anyNonZero = false;
+  for(auto& p : net->parameters()) {
+    if(p->getGrads()) {
+      for(tensorSize_t i = 0; i < p->getGrads()->getSize(); i++) {
+        if((*p->getGrads())[i] != 0.0f) {
+          anyNonZero = true;
+          break;
+        }
+      }
+    }
+  }
+  EXPECT_TRUE(anyNonZero) << "Expected some non-zero gradients before zeroGrad";
+
+  optim->zeroGrad();
+
+  for(auto& p : net->parameters()) {
+    if(p->getGrads()) {
+      for(tensorSize_t i = 0; i < p->getGrads()->getSize(); i++) {
+        ASSERT_NEAR((*p->getGrads())[i], 0.0f, 1e-5)
+          << "Gradient not zeroed at index " << i;
+      }
+    }
+  }
+}
diff --git a/tests/backend/cuda/test_losses.cu b/tests/backend/cuda/test_losses.cu
new file mode 100644
index 0000000..5acc3f2
--- /dev/null
+++ b/tests/backend/cuda/test_losses.cu
@@ -0,0 +1,365 @@
+/**
+ * @file test_losses.cu
+ * @author Robert Baumgartner (r.baumgartner-1@tudelft.nl)
+ * @brief
+ * @version 0.1
+ * @date 2026-05-24
+ *
+ * @copyright Copyright (c) 2026
+ *
+ */
+
+#ifndef __CUDA
+static_assert(false, "File should not be compiled without CUDA enabled");
+#endif // __CUDA
+
+#include <gtest/gtest.h>
+
+#include "data_modeling/tensor_functions.h"
+
+#include "training/loss_functions/rmse_loss.h"
+#include "training/loss_functions/bce_loss.h"
+#include "training/loss_functions/crossentropy_loss.h"
+
+#include <cmath>
+
+using namespace std;
+using namespace train;
+
+TEST(CudaLossTest, CrossEntropyForward) {
+  auto y = TensorFunctions::makeSharedTensor(
+    {2, 3}, {1.0, 0.0, 0.0,
+         0.0, 1.0, 0.0}, Device::CUDA, false);
+
+  auto ypred = TensorFunctions::makeSharedTensor(
+    {2, 3}, {0.7, 0.2, 0.1,
+         0.1, 0.8, 0.1}, Device::CUDA, true);
+
+  CrossEntropyLoss loss;
+  auto result = loss(y, ypred);
+
+  const ftype expected = -(std::log(0.7f) + std::log(0.8f)) / 2.0f;
+  ASSERT_NEAR((*result)[0], expected, 1e-4);
+}
+
+TEST(CudaLossTest, CrossEntropyPerfectPrediction) {
+  auto y = TensorFunctions::makeSharedTensor(
+    {2, 3}, {1.0, 0.0, 0.0,
+         0.0, 1.0, 0.0}, Device::CUDA, false);
+
+  auto ypred = TensorFunctions::makeSharedTensor(
+    {2, 3}, {0.999, 0.0005, 0.0005,
+         0.0005, 0.999, 0.0005}, Device::CUDA, true);
+
+  CrossEntropyLoss loss;
+  auto result = loss(y, ypred);
+
+  ASSERT_LT((*result)[0], 0.01f);
+}
+
+TEST(CudaLossTest, CrossEntropyUniformPrediction) {
+  auto y = TensorFunctions::makeSharedTensor(
+    {1, 3}, {1.0, 0.0, 0.0}, Device::CUDA, false);
+
+  auto ypred = TensorFunctions::makeSharedTensor(
+    {1, 3}, {1.0f/3, 1.0f/3, 1.0f/3}, Device::CUDA, true);
+
+  CrossEntropyLoss loss;
+  auto result = loss(y, ypred);
+
+  ASSERT_NEAR((*result)[0], std::log(3.0f), 1e-4);
+}
+
+TEST(CudaLossTest, CrossEntropyForwardLarge) {
+  constexpr tensorDim_t nSamples = 500;
+  constexpr tensorDim_t nClasses = 200;
+
+  auto yCpu = make_shared<Tensor>(TensorFunctions::Ones({nSamples, nClasses}, Device::CPU, true) * 0.5f);
+  auto yGpu = make_shared<Tensor>(yCpu->createDeepCopy());
+  yGpu->setDevice(Device::CUDA);
+
+  auto ypredCpu = make_shared<Tensor>(TensorFunctions::Ones({nSamples, nClasses}, Device::CPU, true) * 0.7f);
+  auto ypredGpu = make_shared<Tensor>(ypredCpu->createDeepCopy());
+  ypredGpu->setDevice(Device::CUDA);
+
+  CrossEntropyLoss loss;
+  auto resCpu = loss(yCpu, ypredCpu);
+  auto resGpu = loss(yGpu, ypredGpu);
+
+  EXPECT_NEAR((*resCpu)[0], (*resGpu)[0], 0.1);
+}
+
+TEST(CudaLossTest, CrossEntropyThrowsOnDimMismatch) {
+  auto y = TensorFunctions::makeSharedTensor(
+    {2, 3}, {1.0, 0.0, 0.0, 0.0, 1.0, 0.0}, Device::CUDA, false);
+  auto ypred = TensorFunctions::makeSharedTensor(
+    {2, 2}, {0.5, 0.5, 0.5, 0.5}, Device::CUDA, true);
+
+  CrossEntropyLoss loss;
+  ASSERT_THROW(loss(y, ypred), std::invalid_argument);
+}
+
+TEST(CudaLossTest, CrossEntropyBackward) {
+  auto y = TensorFunctions::makeSharedTensor(
+    {2, 3}, {1.0, 0.0, 0.0,
+         0.0, 1.0, 0.0}, Device::CUDA, false);
+  auto ypred = TensorFunctions::makeSharedTensor(
+    {2, 3}, {0.7, 0.2, 0.1,
+         0.1, 0.8, 0.1}, Device::CUDA, true);
+
+  CrossEntropyLoss loss;
+  auto result = loss(y, ypred);
+  result->backward();
+
+  auto grads = ypred->getGrads();
+  ASSERT_NEAR((*grads)[0], -0.7143f, 1e-4);
+  ASSERT_NEAR((*grads)[1],  0.0f,    1e-4);
+  ASSERT_NEAR((*grads)[2],  0.0f,    1e-4);
+  ASSERT_NEAR((*grads)[3],  0.0f,    1e-4);
+  ASSERT_NEAR((*grads)[4], -0.625f,  1e-4);
+  ASSERT_NEAR((*grads)[5],  0.0f,    1e-4);
+}
+
+TEST(CudaLossTest, CrossEntropyBackwardLarge) {
+  constexpr tensorDim_t nSamples = 500;
+  constexpr tensorDim_t nClasses = 200;
+
+  auto yCpu = make_shared<Tensor>(TensorFunctions::Ones({nSamples, nClasses}) * 0.5f);
+  auto yGpu = make_shared<Tensor>(yCpu->createDeepCopy());
+  yGpu->setDevice(Device::CUDA);
+
+  auto ypredCpu = make_shared<Tensor>(TensorFunctions::Ones({nSamples, nClasses}, Device::CPU, true) * 0.7f);
+  auto ypredGpu = make_shared<Tensor>(ypredCpu->createDeepCopy());
+  ypredGpu->setDevice(Device::CUDA);
+
+  CrossEntropyLoss loss;
+  auto resCpu = loss(yCpu, ypredCpu);
+  auto resGpu = loss(yGpu, ypredGpu);
+
+  resCpu->backward();
+  resGpu->backward();
+
+  auto gradsCpu = ypredCpu->getGrads();
+  auto gradsGpu = ypredGpu->getGrads();
+  gradsGpu->setDevice(Device::CPU);
+
+  for(int i = 0; i < ypredCpu->getSize(); i++) {
+    EXPECT_NEAR((*gradsCpu)[i], (*gradsGpu)[i], 1e-4) 
+      << "Failed at index " << i 
+      << " - GradsCpu[i]: " << (*gradsCpu)[i]
+      << " - GradsGpu[i]: " << (*gradsGpu)[i];
+  }
+}
+
+TEST(CudaLossTest, BceForward) {
+  auto y = TensorFunctions::makeSharedTensor(
+    {4, 1}, {0.0, 1.0, 1.0, 0.0}, Device::CUDA, false);
+
+  auto ypred = TensorFunctions::makeSharedTensor(
+    {4, 1}, {0.1, 0.9, 0.8, 0.2}, Device::CUDA, true);
+
+  BceLoss loss;
+  auto result = loss(y, ypred);
+
+  const ftype expected = -(std::log(0.9f) + std::log(0.9f) +
+               std::log(0.8f) + std::log(0.8f)) / 4.0f;
+  ASSERT_NEAR((*result)[0], expected, 1e-4);
+}
+
+TEST(CudaLossTest, BcePerfectPrediction) {
+  auto y = TensorFunctions::makeSharedTensor(
+    {2, 1}, {1.0, 0.0}, Device::CUDA, false);
+
+  auto ypred = TensorFunctions::makeSharedTensor(
+    {2, 1}, {0.999, 0.001}, Device::CUDA, true);
+
+  BceLoss loss;
+  auto result = loss(y, ypred);
+
+  ASSERT_LT((*result)[0], 0.01f);
+}
+
+TEST(CudaLossTest, BceRandomPrediction) {
+  auto y = TensorFunctions::makeSharedTensor(
+    {2, 1}, {1.0, 0.0}, Device::CUDA, false);
+
+  auto ypred = TensorFunctions::makeSharedTensor(
+    {2, 1}, {0.5, 0.5}, Device::CUDA, true);
+
+  BceLoss loss;
+  auto result = loss(y, ypred);
+
+  ASSERT_NEAR((*result)[0], std::log(2.0f), 1e-4);
+}
+
+TEST(CudaLossTest, BceForwardLarge) {
+  constexpr tensorDim_t nSamples = 10000;
+
+  auto yCpu = make_shared<Tensor>(TensorFunctions::Ones({nSamples, 1}) * 0.5f);
+  auto yGpu = make_shared<Tensor>(yCpu->createDeepCopy());
+  yGpu->setDevice(Device::CUDA);
+
+  auto ypredCpu = make_shared<Tensor>(TensorFunctions::Ones({nSamples, 1}) * 0.7f);
+  ypredCpu->setRequiresGrad(true);
+  auto ypredGpu = make_shared<Tensor>(ypredCpu->createDeepCopy());
+  ypredGpu->setDevice(Device::CUDA);
+  ypredGpu->setRequiresGrad(true);
+
+  BceLoss loss;
+  auto resCpu = loss(yCpu, ypredCpu);
+  auto resGpu = loss(yGpu, ypredGpu);
+
+  EXPECT_NEAR((*resCpu)[0], (*resGpu)[0], 1e-4);
+}
+
+TEST(CudaLossTest, BceThrowsOnDimMismatch) {
+  auto y = TensorFunctions::makeSharedTensor(
+    {2, 1}, {1.0, 0.0}, Device::CUDA, false);
+  auto ypred = TensorFunctions::makeSharedTensor(
+    {3, 1}, {0.5, 0.5, 0.5}, Device::CUDA, true);
+
+  BceLoss loss;
+  ASSERT_THROW(loss(y, ypred), std::invalid_argument);
+}
+
+TEST(CudaLossTest, BceNoInfOrNanOnNearZeroPred) {
+  auto y = TensorFunctions::makeSharedTensor(
+    {1, 1}, {1.0}, Device::CUDA, false);
+  auto ypred = TensorFunctions::makeSharedTensor(
+    {1, 1}, {0.0}, Device::CUDA, true);
+
+  BceLoss loss;
+  auto result = loss(y, ypred);
+
+  ASSERT_FALSE(std::isinf((*result)[0]));
+}
+
+TEST(CudaLossTest, BceBackward) {
+  auto y = TensorFunctions::makeSharedTensor(
+    {2, 1}, {1.0, 0.0}, Device::CUDA, false);
+  auto ypred = TensorFunctions::makeSharedTensor(
+    {2, 1}, {0.8, 0.3}, Device::CUDA, true);
+
+  BceLoss loss;
+  auto result = loss(y, ypred);
+  result->backward();
+
+  auto grads = ypred->getGrads();
+  ASSERT_NEAR((*grads)[0], -0.625f,  1e-4);
+  ASSERT_NEAR((*grads)[1],  0.7143f, 1e-4);
+}
+
+TEST(CudaLossTest, BceBackwardLarge) {
+  constexpr tensorDim_t nSamples = 10000;
+
+  auto yCpu = make_shared<Tensor>(TensorFunctions::Ones({nSamples, 1}) * 0.5f);
+  auto yGpu = make_shared<Tensor>(yCpu->createDeepCopy());
+  yGpu->setDevice(Device::CUDA);
+
+  auto ypredCpu = make_shared<Tensor>(TensorFunctions::Ones({nSamples, 1}, Device::CPU, true) * 0.7f);
+  auto ypredGpu = make_shared<Tensor>(ypredCpu->createDeepCopy());
+  ypredGpu->setDevice(Device::CUDA);
+
+  BceLoss loss;
+  auto resCpu = loss(yCpu, ypredCpu);
+  auto resGpu = loss(yGpu, ypredGpu);
+
+  resCpu->backward();
+  resGpu->backward();
+
+  auto gradsCpu = ypredCpu->getGrads();
+  auto gradsGpu = ypredGpu->getGrads();
+  gradsGpu->setDevice(Device::CPU);
+
+  for(int i = 0; i < ypredCpu->getSize(); i++) {
+    EXPECT_NEAR((*gradsCpu)[i], (*gradsGpu)[i], 1e-4) 
+      << "Failed at index " << i 
+      << " - GradsCpu[i]: " << (*gradsCpu)[i]
+      << " - GradsGpu[i]: " << (*gradsGpu)[i];
+  }
+}
+
+TEST(CudaLossTest, RmseForward) {
+  auto y = TensorFunctions::makeSharedTensor(
+    {3}, {1.0, 2.0, 3.0}, Device::CUDA, false);
+  auto ypred = TensorFunctions::makeSharedTensor(
+    {3}, {1.5, 2.5, 2.5}, Device::CUDA, true);
+
+  RmseLoss loss;
+  auto result = loss(y, ypred);
+
+  ASSERT_NEAR((*result)[0], 0.5f, 1e-4);
+}
+
+TEST(CudaLossTest, RmseForwardLarge) {
+  auto yCpu = make_shared<Tensor>(TensorFunctions::Gaussian({500, 500}, 1.0f));
+  auto yGpu = make_shared<Tensor>(yCpu->createDeepCopy());
+  yGpu->setDevice(Device::CUDA);
+
+  auto ypredCpu = make_shared<Tensor>(TensorFunctions::Gaussian({500, 500}, 1.0f));
+  ypredCpu->setRequiresGrad(true);
+  auto ypredGpu = make_shared<Tensor>(ypredCpu->createDeepCopy());
+  ypredGpu->setDevice(Device::CUDA);
+  ypredGpu->setRequiresGrad(true);
+
+  RmseLoss loss;
+  auto resCpu = loss(yCpu, ypredCpu);
+  auto resGpu = loss(yGpu, ypredGpu);
+
+  EXPECT_NEAR((*resCpu)[0], (*resGpu)[0], 0.1);
+}
+
+TEST(CudaLossTest, RmsePerfectPrediction) {
+  auto y = TensorFunctions::makeSharedTensor(
+    {3}, {1.0, 2.0, 3.0}, Device::CUDA, false);
+  auto ypred = TensorFunctions::makeSharedTensor(
+    {3}, {1.0, 2.0, 3.0}, Device::CUDA, true);
+
+  RmseLoss loss;
+  auto result = loss(y, ypred);
+
+  ASSERT_NEAR((*result)[0], 0.0f, 1e-4);
+}
+
+TEST(CudaLossTest, RmseBackward) {
+  auto y = TensorFunctions::makeSharedTensor(
+    {2}, {1.0, 0.0}, Device::CUDA, false);
+  auto ypred = TensorFunctions::makeSharedTensor(
+    {2}, {0.5, 0.5}, Device::CUDA, true);
+
+  RmseLoss loss;
+  auto result = loss(y, ypred);
+  result->backward();
+
+  auto grads = ypred->getGrads();
+  ASSERT_NEAR((*grads)[0], -0.5f, 1e-4);
+  ASSERT_NEAR((*grads)[1],  0.5f, 1e-4);
+}
+
+TEST(CudaLossTest, RmseBackwardLarge) {
+  auto yCpu = make_shared<Tensor>(TensorFunctions::Gaussian({500, 500}, 1.0f));
+  auto yGpu = make_shared<Tensor>(yCpu->createDeepCopy());
+  yGpu->setDevice(Device::CUDA);
+
+  auto ypredCpu = make_shared<Tensor>(TensorFunctions::Gaussian({500, 500}, 1.0f, true));
+  auto ypredGpu = make_shared<Tensor>(ypredCpu->createDeepCopy());
+  ypredGpu->setDevice(Device::CUDA);
+
+  RmseLoss loss;
+  auto resCpu = loss(yCpu, ypredCpu);
+  auto resGpu = loss(yGpu, ypredGpu);
+
+  resCpu->backward();
+  resGpu->backward();
+
+  auto gradsCpu = ypredCpu->getGrads();
+  auto gradsGpu = ypredGpu->getGrads();
+  gradsGpu->setDevice(Device::CPU);
+
+  for(int i = 0; i < ypredCpu->getSize(); i++) {
+    EXPECT_NEAR((*gradsCpu)[i], (*gradsGpu)[i], 1e-4) 
+      << "Failed at index " << i 
+      << " - GradsCpu[i]: " << (*gradsCpu)[i]
+      << " - GradsGpu[i]: " << (*gradsGpu)[i];
+  }
+}
diff --git a/tests/backend/cuda/test_module_cuda.cu b/tests/backend/cuda/test_module_cuda.cu
new file mode 100644
index 0000000..7626e21
--- /dev/null
+++ b/tests/backend/cuda/test_module_cuda.cu
@@ -0,0 +1,479 @@
+/**
+ * @file test_module_cuda.cu
+ * @author Robert Baumgartner (r.baumgartner-1@tudelft.nl)
+ * @brief
+ * @version 0.1
+ * @date 2026-05-13
+ *
+ * @copyright Copyright (c) 2026
+ *
+ */
+
+#ifndef __CUDA
+static_assert(false, "File should not be compiled without CUDA enabled");
+#endif // __CUDA
+
+#include <gtest/gtest.h>
+
+#include "data_modeling/tensor_functions.h"
+#include "computational_graph/tensor_ops/graph_creation.h"
+
+#include "module/activation_functions/relu.h"
+#include "module/activation_functions/leaky_relu.h"
+#include "module/activation_functions/softmax.h"
+#include "module/activation_functions/sigmoid.h"
+#include "module/layers/ff_layer.h"
+
+using namespace std;
+
+TEST(CudaActivationTest, ReluForward) {
+  auto t1 = TensorFunctions::Ones({300, 500}, Device::CUDA);
+  auto f = module::ReLu();
+
+  auto res = f(t1);
+  res.setDevice(Device::CPU);
+  t1.setDevice(Device::CPU);
+
+  for(size_t i=0; i<t1.getSize(); i++){
+    ASSERT_NEAR(res[i], t1[i], 1e-5);
+  }
+}
+
+TEST(CudaActivationTest, ReluInputNegative) {
+  auto t1 = TensorFunctions::Ones({300, 500}, Device::CUDA) * -1;
+  auto f = module::ReLu();
+
+  auto res = f(t1);
+  res.setDevice(Device::CPU);
+
+  constexpr ftype zero = 0;
+  for(size_t i=0; i<t1.getSize(); i++){
+    ASSERT_DOUBLE_EQ(res[i], zero);
+  }
+}
+
+TEST(CudaAutogradTest, ReLUBackward) {
+  auto x = TensorFunctions::makeSharedTensor({3}, {-1.0, 0.0, 2.0}, Device::CUDA, true);
+  auto relu = module::ReLu();
+
+  auto y = relu(x);
+  auto loss = cgraph::sumTensor(y);
+
+  loss->backward();
+
+  ASSERT_NEAR(x->getGrads()->get(0), 0.0, 1e-5);
+  ASSERT_NEAR(x->getGrads()->get(1), 0.0, 1e-5);
+  ASSERT_NEAR(x->getGrads()->get(2), 1.0, 1e-5);
+}
+
+TEST(CudaActivationTest, LeakyReluForward) {
+  auto t1 = TensorFunctions::Ones({300, 500}, Device::CUDA);
+
+  auto f = module::LeakyReLu(0.3);
+  auto res = f(t1);
+
+  res.setDevice(Device::CPU);
+  t1.setDevice(Device::CPU);
+
+  for(size_t i=0; i<t1.getSize(); i++){
+    ASSERT_NEAR(res[i], t1[i], 1e-5);
+  }
+}
+
+TEST(CudaActivationTest, LeakyReluInputNegative) {
+  constexpr ftype factor = -1;
+  auto t1 = TensorFunctions::Ones({300, 500}, Device::CUDA) * factor;
+
+  constexpr ftype eps = 0.3;
+  auto f = module::LeakyReLu(eps);
+  auto res = f(t1);
+
+  res.setDevice(Device::CPU);
+
+  for(size_t i=0; i<t1.getSize(); i++){
+    ASSERT_NEAR(res[i], factor * eps, 1e-5);
+  }
+}
+
+TEST(CudaAutogradTest, LeakyReLUBackward) {
+  auto x = TensorFunctions::makeSharedTensor({3}, {-1.0, 0.0, 2.0}, Device::CUDA, true);
+
+  constexpr ftype eps = 0.3;
+  auto relu = module::LeakyReLu(eps);
+
+  auto y = relu(x);
+  auto loss = cgraph::sumTensor(y);
+
+  loss->backward();
+
+  ASSERT_NEAR(x->getGrads()->get(0), eps, 1e-5);
+  ASSERT_NEAR(x->getGrads()->get(1), eps, 1e-5);
+  ASSERT_NEAR(x->getGrads()->get(2), 1.0, 1e-5);
+}
+
+TEST(CudaActivationTest, SigmoidLargePositive) {
+  auto t = TensorFunctions::makeSharedTensor({1}, vector<ftype>{100.0}, Device::CUDA);
+
+  module::Sigmoid sig;
+  auto res = sig(*t);
+
+  ASSERT_NEAR(res[0], 1.0, 1e-5);
+  ASSERT_FALSE(std::isnan(res[0]));
+  ASSERT_FALSE(std::isinf(res[0]));
+}
+
+TEST(CudaActivationTest, SigmoidLargeNegative) {
+  auto t = TensorFunctions::makeSharedTensor({1}, vector<ftype>{-100.0}, Device::CUDA);
+
+  module::Sigmoid sig;
+  auto res = sig(*t);
+
+  ASSERT_NEAR(res[0], 0.0, 1e-5);
+  ASSERT_FALSE(std::isnan(res[0]));
+  ASSERT_FALSE(std::isinf(res[0]));
+}
+
+TEST(CudaAutogradTest, SigmoidBackward) {
+  auto t = TensorFunctions::makeSharedTensor({2}, {0.0, 1.0}, Device::CUDA, true);
+
+  module::Sigmoid sig;
+  auto res = sig(t);
+  res->backward();
+
+  auto grads = t->getGrads();
+  ASSERT_NEAR((*grads)[0], 0.25, 1e-4);
+  ASSERT_NEAR((*grads)[1], 0.1966, 1e-4);
+}
+
+TEST(CudaActivationTest, SoftmaxForward) {
+  auto t = TensorFunctions::makeSharedTensor({1, 3}, {1.0, 2.0, 3.0}, Device::CUDA);
+
+  module::Softmax sm;
+  auto res = sm(*t);
+
+  ASSERT_NEAR(res[0], 0.0900, 1e-4);
+  ASSERT_NEAR(res[1], 0.2447, 1e-4);
+  ASSERT_NEAR(res[2], 0.6652, 1e-4);
+}
+
+TEST(CudaActivationTest, SoftmaxSumsToOne) {
+  auto t = TensorFunctions::makeSharedTensor({2, 4},
+                  {1.0, 2.0, 3.0, 4.0,
+                   2.0, 1.0, 4.0, 3.0}, Device::CUDA);
+
+  module::Softmax sm;
+  auto res = sm(*t);
+
+  ftype row0sum = res[0] + res[1] + res[2] + res[3];
+  ftype row1sum = res[4] + res[5] + res[6] + res[7];
+  ASSERT_NEAR(row0sum, 1.0, 1e-5);
+  ASSERT_NEAR(row1sum, 1.0, 1e-5);
+}
+
+TEST(CudaActivationTest, SoftmaxForwardNumericalStability) {
+  auto t = TensorFunctions::makeSharedTensor({1, 3}, {100.0, 101.0, 102.0}, Device::CUDA);
+
+  module::Softmax sm;
+  auto res = sm(*t);
+
+  for(int i = 0; i < 3; i++) {
+    ASSERT_FALSE(std::isnan(res[i]));
+    ASSERT_FALSE(std::isinf(res[i]));
+  }
+
+  ftype rowsum = res[0] + res[1] + res[2];
+  ASSERT_NEAR(rowsum, 1.0, 1e-5);
+}
+
+TEST(CudaActivationTest, SoftmaxMediumLarge) {
+  constexpr tensorDim_t testDim = 190;
+  assert(testDim <= 256 && testDim > 64); // see the kernel call 
+
+  auto t = TensorFunctions::Gaussian({5, 10, testDim}, 2.0f, Device::CUDA);
+  auto tCopy = t.createDeepCopy();
+  tCopy.setDevice(Device::CPU);
+
+  module::Softmax sm;
+  auto resGpu = sm(t);
+  auto resCpu = sm(tCopy);
+
+  resGpu.setDevice(Device::CUDA);
+  for(int i = 0; i < resGpu.getDims().get(0); i++) {
+    for(int j = 0; j < resGpu.getDims().get(1); j++) {
+      for(int k = 0; k < resGpu.getDims().get(2); k++) {
+        ASSERT_NEAR(resCpu.get(i, j, k), resGpu.get(i, j, k), 1e-4)
+          << "Mismatch at (" << i << ", " << j << ", " << k <<  ")"
+          << " cpu=" << resCpu.get(i, j, k) 
+          << " gpu=" << resGpu.get(i, j, k);
+      }
+    }
+  }
+}
+
+TEST(CudaActivationTest, SoftmaxLarge) {
+  constexpr tensorDim_t testDim = 1500;
+  assert(testDim > 256); // see the kernel call 
+
+  auto t = TensorFunctions::Gaussian({2, 2, testDim}, 2.0f, Device::CUDA);
+  auto tCopy = t.createDeepCopy();
+  tCopy.setDevice(Device::CPU);
+
+  module::Softmax sm;
+  auto resGpu = sm(t);
+  auto resCpu = sm(tCopy);
+
+  resGpu.setDevice(Device::CUDA);
+  for(int i = 0; i < resGpu.getDims().get(0); i++) {
+    for(int j = 0; j < resGpu.getDims().get(1); j++) {
+      for(int k = 0; k < resGpu.getDims().get(2); k++) {
+        ASSERT_NEAR(resCpu.get(i, j, k), resGpu.get(i, j, k), 1e-4)
+          << "Mismatch at (" << i << ", " << j << ", " << k <<  ")"
+          << " cpu=" << resCpu.get(i, j, k) 
+          << " gpu=" << resGpu.get(i, j, k);
+      }
+    }
+  }
+}
+
+TEST(CudaAutogradTest, SoftmaxBackward) {
+  auto t = TensorFunctions::makeSharedTensor({1, 3}, {1.0, 2.0, 3.0}, Device::CUDA, true);
+
+  module::Softmax sm;
+  auto resPtr = sm(t);
+
+  auto upstream = TensorFunctions::makeSharedTensor({1, 3}, {1.0, 0.0, 0.0}, Device::CUDA);
+  resPtr->setGrads(upstream);
+  resPtr->backward();
+
+  auto grads = t->getGrads();
+  ASSERT_NEAR((*grads)[0],  0.0819, 1e-4);
+  ASSERT_NEAR((*grads)[1], -0.0220, 1e-4);
+  ASSERT_NEAR((*grads)[2], -0.0599, 1e-4);
+}
+
+TEST(CudaAutogradTest, SoftmaxBackwardBatched) {
+  // 4 samples x 3 classes — exercises the one-block path with multiple strides
+  auto tCpu = make_shared<Tensor>(TensorFunctions::Gaussian({4, 3}, 1.0f, true));
+  auto tGpu = make_shared<Tensor>(tCpu->createDeepCopy());
+  tGpu->setDevice(Device::CUDA);
+
+  module::Softmax sm;
+  auto resCpu = sm(tCpu);
+  auto resGpu = sm(tGpu);
+
+  resCpu->backward();
+  resGpu->backward();
+
+  auto gradsCpu = tCpu->getGrads();
+  auto gradsGpu = tGpu->getGrads();
+
+  for(int i = 0; i < tCpu->getSize(); i++) {
+    EXPECT_NEAR((*gradsCpu)[i], (*gradsGpu)[i], 1e-4) 
+      << "Failed at index " << i 
+      << " - GradsCpu[i]: " << (*gradsCpu)[i]
+      << " - GradsGpu[i]: " << (*gradsGpu)[i];
+  }
+}
+
+TEST(CudaAutogradTest, SoftmaxBackwardLarge) {
+  constexpr tensorDim_t testDim = 1300;
+  auto tCpu = make_shared<Tensor>(TensorFunctions::Gaussian({2, 2, testDim}, 2.0f, true));
+  auto tGpu = make_shared<Tensor>(tCpu->createDeepCopy());
+  tGpu->setDevice(Device::CUDA);
+
+  module::Softmax sm;
+  auto resPtrCpu = sm(tCpu);
+  auto resPtrGpu = sm(tGpu);
+
+  resPtrCpu->backward();
+  resPtrGpu->backward();
+
+  auto gradsCpu = tCpu->getGrads();
+  auto gradsGpu = tGpu->getGrads();
+
+  for(int i = 0; i < tCpu->getSize(); i++) {
+    EXPECT_NEAR((*gradsCpu)[i], (*gradsGpu)[i], 1e-4) 
+      << "Failed at index " << i 
+      << " - GradsCpu[i]: " << (*gradsCpu)[i]
+      << " - GradsGpu[i]: " << (*gradsGpu)[i];
+  }
+}
+
+TEST(CudaLayerTest, TestFfLayer) {
+  auto t1 = TensorFunctions::Ones({3, 2}, Device::CUDA);
+  auto layer = module::FfLayer(2, 1, Device::CUDA);
+
+  auto res = layer(t1);
+
+  ASSERT_EQ(res.getDims(), Dimension({3, 1}));
+}
+
+TEST(CudaLayerTest, TestFfLayerLarge) {
+  constexpr tensorDim_t largeDim = 200;
+  auto tCpu = TensorFunctions::Ones({30, largeDim});
+
+  auto tGpu = tCpu.createDeepCopy();
+  tGpu.setDevice(Device::CUDA);
+
+  auto layer = module::FfLayer(largeDim, 10);
+  auto resCpu = layer(tCpu);
+
+  layer.setDevice(Device::CUDA);
+  auto resGpu = layer(tGpu);
+
+  for(int i = 0; i < resCpu.getSize(); i++) {
+    EXPECT_NEAR(resCpu[i], resGpu[i], 1e-4);
+  }
+}
+
+TEST(CudaAutogradTest, FfLayerBackward) {
+  auto x = TensorFunctions::makeSharedTensor({2, 3}, {
+    1.0, 1.0, 1.0,
+    1.0, 1.0, 1.0
+  }, Device::CUDA, true);
+
+  auto layer = module::FfLayer(3, 2, Device::CUDA, false, true);
+  auto w = layer.getWeights();
+  for(tensorSize_t i = 0; i < w->getSize(); i++) w->set(1.0, i);
+
+  auto res = layer(x);
+  auto loss = cgraph::sumTensor(res);
+  loss->backward();
+
+  auto xGrads = x->getGrads();
+  ASSERT_NE(xGrads, nullptr);
+  ASSERT_EQ(xGrads->getDims(), x->getDims());
+  for(tensorSize_t i = 0; i < xGrads->getSize(); i++) {
+    ASSERT_NEAR((*xGrads)[i], 2.0, 1e-5);
+  }
+
+  auto wGrads = layer.getWeights()->getGrads();
+  ASSERT_NE(wGrads, nullptr);
+  ASSERT_EQ(wGrads->getDims(), layer.getWeights()->getDims());
+  for(tensorSize_t i = 0; i < wGrads->getSize(); i++) {
+    ASSERT_NEAR((*wGrads)[i], 2.0, 1e-5);
+  }
+}
+
+TEST(CudaAutogradTest, FfLayerBackwardWithBias) {
+  auto x = TensorFunctions::makeSharedTensor({2, 3}, {
+    1.0, 1.0, 1.0,
+    1.0, 1.0, 1.0
+  }, Device::CUDA, true);
+
+  auto layer = module::FfLayer(3, 2, Device::CUDA, true, true);
+  auto w = layer.getWeights();
+  for(tensorSize_t i = 0; i < w->getSize(); i++) w->set(1.0, i);
+
+  auto res = layer(x);
+  auto loss = cgraph::sumTensor(res);
+  loss->backward();
+
+  auto xGrads = x->getGrads();
+  ASSERT_NE(xGrads, nullptr);
+  for(tensorSize_t i = 0; i < xGrads->getSize(); i++) {
+    ASSERT_NEAR((*xGrads)[i], 2.0, 1e-5);
+  }
+
+  auto wGrads = layer.getWeights()->getGrads();
+  ASSERT_NE(wGrads, nullptr);
+  for(tensorSize_t i = 0; i < wGrads->getSize(); i++) {
+    ASSERT_NEAR((*wGrads)[i], 2.0, 1e-5);
+  }
+
+  auto bGrads = layer.getBias()->getGrads();
+  ASSERT_NE(bGrads, nullptr);
+  ASSERT_NEAR((*bGrads)[0], 2.0, 1e-5);
+  ASSERT_NEAR((*bGrads)[1], 2.0, 1e-5);
+}
+
+TEST(CudaAutogradTest, FfLayerBackwardLarge) {
+  constexpr tensorDim_t inDim = 200;
+  constexpr tensorDim_t outDim = 10;
+
+  auto xCpu = make_shared<Tensor>(TensorFunctions::Gaussian({30, inDim}, 1.0f, true));
+  auto xGpu = make_shared<Tensor>(xCpu->createDeepCopy());
+  xGpu->setDevice(Device::CUDA);
+
+  auto layerCpu = module::FfLayer(inDim, outDim, Device::CPU, false, true);
+  auto layerGpu = module::FfLayer(inDim, outDim, Device::CUDA, false, true);
+
+  // give both layers identical weights
+  auto wCpu = layerCpu.getWeights();
+  auto wGpu = layerGpu.getWeights();
+  for(tensorSize_t i = 0; i < wCpu->getSize(); i++) {
+    const ftype val = 1.0f / inDim;
+    wCpu->set(val, i);
+    wGpu->set(val, i);
+  }
+
+  auto resCpu = layerCpu(xCpu);
+  auto resGpu = layerGpu(xGpu);
+
+  resCpu->backward();
+  resGpu->backward();
+
+  auto xGradsCpu = xCpu->getGrads();
+  auto xGradsGpu = xGpu->getGrads();
+  for(int i = 0; i < xCpu->getSize(); i++) {
+    EXPECT_NEAR((*xGradsCpu)[i], (*xGradsGpu)[i], 1e-4);
+  }
+
+  auto wGradsCpu = layerCpu.getWeights()->getGrads();
+  auto wGradsGpu = layerGpu.getWeights()->getGrads();
+  for(int i = 0; i < wCpu->getSize(); i++) {
+    EXPECT_NEAR((*wGradsCpu)[i], (*wGradsGpu)[i], 1e-4);
+  }
+}
+
+TEST(CudaAutogradTest, FfLayerBackwardWithBiasLarge) {
+  constexpr tensorDim_t inDim = 200;
+  constexpr tensorDim_t outDim = 10;
+
+  auto xCpu = make_shared<Tensor>(TensorFunctions::Gaussian({30, inDim}, 1.0f, true));
+  auto xGpu = make_shared<Tensor>(xCpu->createDeepCopy());
+  xGpu->setDevice(Device::CUDA);
+
+  auto layerCpu = module::FfLayer(inDim, outDim, Device::CPU, true, true);
+  auto layerGpu = module::FfLayer(inDim, outDim, Device::CUDA, true, true);
+
+  // give both layers identical weights and biases
+  auto wCpu = layerCpu.getWeights();
+  auto wGpu = layerGpu.getWeights();
+  for(tensorSize_t i = 0; i < wCpu->getSize(); i++) {
+    const ftype val = 1.0f / inDim;
+    wCpu->set(val, i);
+    wGpu->set(val, i);
+  }
+
+  auto bCpu = layerCpu.getBias();
+  auto bGpu = layerGpu.getBias();
+  for(tensorSize_t i = 0; i < bCpu->getSize(); i++) {
+    bCpu->set(0.1f, i);
+    bGpu->set(0.1f, i);
+  }
+
+  auto resCpu = layerCpu(xCpu);
+  auto resGpu = layerGpu(xGpu);
+
+  resCpu->backward();
+  resGpu->backward();
+
+  auto xGradsCpu = xCpu->getGrads();
+  auto xGradsGpu = xGpu->getGrads();
+  for(int i = 0; i < xCpu->getSize(); i++) {
+    EXPECT_NEAR((*xGradsCpu)[i], (*xGradsGpu)[i], 1e-4);
+  }
+
+  auto wGradsCpu = layerCpu.getWeights()->getGrads();
+  auto wGradsGpu = layerGpu.getWeights()->getGrads();
+  for(int i = 0; i < wCpu->getSize(); i++) {
+    EXPECT_NEAR((*wGradsCpu)[i], (*wGradsGpu)[i], 1e-4);
+  }
+
+  auto bGradsCpu = layerCpu.getBias()->getGrads();
+  auto bGradsGpu = layerGpu.getBias()->getGrads();
+  for(int i = 0; i < bCpu->getSize(); i++) {
+    EXPECT_NEAR((*bGradsCpu)[i], (*bGradsGpu)[i], 1e-4);
+  }
+}
\ No newline at end of file
diff --git a/tests/backend/cuda/test_tensorops_cuda.cu b/tests/backend/cuda/test_tensorops_cuda.cu
new file mode 100644
index 0000000..8178d72
--- /dev/null
+++ b/tests/backend/cuda/test_tensorops_cuda.cu
@@ -0,0 +1,443 @@
+/**
+ * @file test_tensorops_cuda.cu
+ * @author Robert Baumgartner (r.baumgartner-1@tudelft.nl)
+ * @brief 
+ * @version 0.1
+ * @date 2026-05-12
+ * 
+ * @copyright Copyright (c) 2026
+ * 
+ */
+
+#ifndef __CUDA
+static_assert(false, "File should not be compiled without CUDA enabled");
+#endif // __CUDA
+
+#include <gtest/gtest.h>
+
+#include "data_modeling/tensor.h"
+#include "data_modeling/tensor_functions.h"
+
+#include "computational_graph/tensor_ops/graph_creation.h"
+
+TEST(CudaTensorOpsTest, TestCtor) {
+  auto t = Tensor({2, 2}, {2.0, 3.0, 4.0, 5.0}, Device::CUDA);
+
+  ASSERT_EQ(t.getDims(), Dimension({2, 2}));
+  ASSERT_EQ(t.getDevice(), Device::CUDA);
+  ASSERT_TRUE(!t.getRequiresGrad());
+
+  ASSERT_NEAR(t.get(0, 0), 2.0, 1e-5);
+  ASSERT_NEAR(t.get(0, 1), 3.0, 1e-5);
+  ASSERT_NEAR(t.get(1, 0), 4.0, 1e-5);
+  ASSERT_NEAR(t.get(1, 1), 5.0, 1e-5);
+}
+
+TEST(CudaTensorOpsTest, ScalarAdd) {
+  auto t1 = TensorFunctions::Ones({500, 500}, Device::CUDA);
+
+  auto res = t1 + 1.5;
+  res.setDevice(Device::CPU);
+
+  constexpr ftype sum = 2.5;
+  for(auto i = 0; i < t1.getDims().get(0); i++) {
+    for(auto j = 0; j < t1.getDims().get(1); j++) {
+      ASSERT_NEAR(res.get(i, j), sum, 1e-5);
+    }
+  }
+}
+
+TEST(CudaTensorOpsTest, ScalarMul) {
+  auto t1 = TensorFunctions::Ones({500, 500}, Device::CUDA);
+
+  constexpr ftype f = 2.5;
+  auto res = t1 * f;
+  res.setDevice(Device::CPU);
+
+  for(auto i=0; i<t1.getDims().get(0); i++) {
+    for(auto j=0; j<t1.getDims().get(1); j++) {
+      ASSERT_NEAR(res.get(i, j), f, 1e-5);
+    }
+  }
+}
+
+TEST(CudaAutogradTest, ScalarMul) {
+  auto t1 = TensorFunctions::makeSharedTensor({1}, {2.0}, Device::CUDA, true);
+  auto t2 = TensorFunctions::makeSharedTensor({1}, {3.0}, Device::CUDA, true);
+
+  auto t3 = cgraph::mul(t1, t2);
+  auto loss = cgraph::mul(t3, t3);
+
+  loss->backward();
+
+  ASSERT_DOUBLE_EQ(t1->getGrads()->get(0), 36.0);
+  ASSERT_DOUBLE_EQ(t2->getGrads()->get(0), 24.0);
+}
+
+TEST(CudaTensorOpsTest, TensorAdd) {
+  auto t1 = TensorFunctions::Ones({500, 500}, Device::CUDA);
+  auto t2 = TensorFunctions::Ones({500, 500}, Device::CUDA) * 4;
+
+  auto res = t1 + t2;
+  res.setDevice(Device::CPU);
+
+  constexpr ftype sum = 5.0;
+  for(auto i=0; i<t1.getDims().get(0); i++) {
+    for(auto j=0; j<t1.getDims().get(1); j++) {
+      ASSERT_NEAR(res.get(i, j), sum, 1e-5);
+    }
+  }
+}
+
+TEST(CudaTensorOpsTest, TensorAddThrowsOnDimMismatch) {
+  auto t1 = TensorFunctions::Ones({2, 2}, Device::CUDA);
+  auto t2 = TensorFunctions::Ones({2, 3}, Device::CUDA) * 4;
+
+  ASSERT_THROW(t1 + t2, std::invalid_argument);
+}
+
+TEST(CudaAutogradTest, TensorAdd) {
+  auto t1 = TensorFunctions::makeSharedTensor({1}, {3.0}, Device::CUDA, true);
+  auto t2 = TensorFunctions::makeSharedTensor({1}, {2.0}, Device::CUDA, true);
+
+  auto t3 = cgraph::add(t1, t2);
+  auto loss = cgraph::mul(t3, t3);
+
+  loss->backward();
+
+  ASSERT_NEAR(t1->getGrads()->get(0), 10.0, 1e-5);
+  ASSERT_NEAR(t2->getGrads()->get(0), 10.0, 1e-5);
+}
+
+TEST(CudaTensorOpsTest, BroadcastAdd) {
+  auto t1 = TensorFunctions::Ones({3, 2, 2}, Device::CUDA);
+  auto t2 = Tensor({2}, {2, 3}, Device::CUDA);
+
+  auto res = t1 + t2;
+
+  ASSERT_EQ(res.getDims(), t1.getDims());
+  
+  for(auto i=0; i<res.getDims().get(0); i++) {
+    for(auto j=0; j<res.getDims().get(1); j++) {
+      ASSERT_NEAR(res.get(i, j, 0), 3.0, 1e-5);
+      ASSERT_NEAR(res.get(i, j, 1), 4.0, 1e-5);
+    }
+  }
+}
+
+TEST(CudaTensorOpsTest, BroadcastAdd_2D) {
+  // (2,3) + (3)
+  auto t1 = Tensor({2, 3}, {1.0, 2.0, 3.0,
+                            4.0, 5.0, 6.0}, Device::CUDA);
+  auto t2 = Tensor({3}, {10.0, 20.0, 30.0}, Device::CUDA);
+
+  auto res = t1 + t2;
+
+  // expected: each row of t1 gets t2 added elementwise
+  ASSERT_NEAR(res.get(0, 0), 11.0, 1e-5);
+  ASSERT_NEAR(res.get(0, 1), 22.0, 1e-5);
+  ASSERT_NEAR(res.get(0, 2), 33.0, 1e-5);
+  ASSERT_NEAR(res.get(1, 0), 14.0, 1e-5);
+  ASSERT_NEAR(res.get(1, 1), 25.0, 1e-5);
+  ASSERT_NEAR(res.get(1, 2), 36.0, 1e-5);
+}
+
+TEST(CudaTensorOpsTest, BroadcastAddNotCommutative) {
+  auto t1 = TensorFunctions::Ones({3, 2, 2}, Device::CUDA);
+  auto t2 = Tensor({2}, {2, 3}, Device::CUDA);
+
+  ASSERT_THROW(t2 + t1, std::invalid_argument);
+}
+
+TEST(CudaAutogradTest, BroadcastAdd) {
+  auto t1 = TensorFunctions::makeSharedTensor({2, 3},
+    {1.0, 2.0, 3.0,
+     4.0, 5.0, 6.0}, Device::CUDA, true);
+  auto bias = TensorFunctions::makeSharedTensor({3},
+    {0.0, 0.0, 0.0}, Device::CUDA, true);
+
+  auto res = cgraph::add(t1, bias);
+  res->backward();
+
+  // bias grad should be sum over batch: [2, 2, 2]
+  auto biasGrad = bias->getGrads();
+  ASSERT_NEAR((*biasGrad)[0], 2.0, 1e-5);
+  ASSERT_NEAR((*biasGrad)[1], 2.0, 1e-5);
+  ASSERT_NEAR((*biasGrad)[2], 2.0, 1e-5);
+
+  // t1 grad should be ones (add is identity for non-broadcast operand)
+  auto t1Grad = t1->getGrads();
+  for(int i = 0; i < 6; i++) {
+    ASSERT_NEAR((*t1Grad)[i], 1.0, 1e-5);
+  }
+}
+
+TEST(CudaAutogradTest, BroadcastAddLarge) {
+  constexpr int dimsize = 1500;
+
+  auto t1 = std::make_shared<Tensor>(
+    TensorFunctions::Gaussian({2, dimsize},
+      5.0, Device::CPU, true));
+
+  auto bias = std::make_shared<Tensor>(
+    TensorFunctions::Gaussian({dimsize},
+      2.0, Device::CPU, true));
+
+  auto t1Gpu = std::make_shared<Tensor>(t1->createDeepCopy());
+  t1Gpu->setDevice(Device::CUDA);
+  auto biasGpu = std::make_shared<Tensor>(bias->createDeepCopy());
+  biasGpu->setDevice(Device::CUDA);
+
+  auto resCpu = cgraph::add(t1, bias);
+  resCpu->backward();
+
+  auto resGpu = cgraph::add(t1Gpu, biasGpu);
+  resGpu->backward();
+  resGpu->setDevice(Device::CPU);
+
+  auto biasGrad = bias->getGrads();
+  auto biasGradGpu = biasGpu->getGrads();
+  for(int i = 0; i < biasGrad->getSize(); i++) {
+    ASSERT_NEAR((*biasGrad)[i], (*biasGradGpu)[i], 1e-5);
+  }
+
+  auto t1Grad = t1->getGrads();
+  auto t1GpuGrads = t1Gpu->getGrads();
+  for(int i = 0; i < t1Grad->getSize(); i++) {
+    ASSERT_NEAR((*t1Grad)[i], (*t1GpuGrads)[i], 1e-5);
+  }
+}
+
+TEST(CudaTensorOpsTest, MatrixAdd) {
+  auto t1 = TensorFunctions::Ones({200, 200}, Device::CUDA);
+  auto t2 = TensorFunctions::Ones({200, 200}, Device::CUDA);
+    
+  auto res = t1 + t2;
+  res.setDevice(Device::CPU);
+
+  constexpr ftype resSum = 2.0;
+  for(auto i=0; i<t1.getDims().get(0); i++) {
+    for(auto j=0; j<t1.getDims().get(1); j++) {
+      ASSERT_NEAR(res.get(i, j), resSum, 1e-5);
+    }
+  }
+}
+
+TEST(CudaTensorOpsTest, ElementwiseMul) {
+  constexpr ftype factor = 0.5;
+  auto t1 = TensorFunctions::Ones({200, 200}, Device::CUDA);
+  auto t2 = TensorFunctions::Ones({200, 200}, Device::CUDA) * 0.5;
+    
+  auto res = t1 * t2;
+  res.setDevice(Device::CPU);
+
+  for(auto i=0; i<t1.getDims().get(0); i++) {
+    for(auto j=0; j<t1.getDims().get(1); j++) {
+      ASSERT_NEAR(res.get(i, j), factor, 1e-5);
+    }
+  }
+}
+
+TEST(CudaTensorOpsTest, ElementwiseMulThrowsOnDimensionMismatch) {
+  auto t1 = TensorFunctions::Ones({2, 2}, Device::CUDA);
+  auto t2 = TensorFunctions::Ones({2, 3}, Device::CUDA) * 0.5;
+    
+  ASSERT_THROW(t1 * t2, std::invalid_argument);
+}
+
+TEST(CudaTensorOpsTest, MatMul) {
+  auto t1 = Tensor({2, 2});
+  auto t2 = Tensor({2, 2});
+
+  auto cmpRes = Tensor({2, 2});
+
+  auto populateTensor = [](Tensor& t, ftype v1, ftype v2, ftype v3, ftype v4) {
+    t.set(v1, {0, 0});
+    t.set(v2, {0, 1});
+    t.set(v3, {1, 0});
+    t.set(v4, {1, 1});
+  };
+
+  populateTensor(t1, 1, 2, 3, 4);
+  populateTensor(t2, 5, 6, 7, 8);
+  populateTensor(cmpRes, 19, 22, 43, 50);
+
+  t1.setDevice(Device::CUDA);
+  t2.setDevice(Device::CUDA);
+
+  auto res = t1.matmul(t2);
+  res.setDevice(Device::CPU);
+
+  auto expectedDims = std::vector<tensorDim_t>{2, 2};
+  ASSERT_EQ(res.getDims().toVector(), expectedDims);
+
+  for(auto i=0; i<t1.getDims().get(0); i++) {
+    for(auto j=0; j<t1.getDims().get(1); j++) {
+      EXPECT_NEAR(res.get(i, j), cmpRes.get(i, j), 1e-5);
+    }
+  }
+}
+
+TEST(CudaTensorOpsTest, MatMulLarge) {
+  auto t1 = TensorFunctions::Gaussian({1000, 10}, 2.0);
+  auto t2 = TensorFunctions::Gaussian({10, 1000}, 2.0);
+  auto resCpu = t1.matmul(t2);
+
+  t1.setDevice(Device::CUDA);
+  t2.setDevice(Device::CUDA);
+
+  auto resGpu = t1.matmul(t2);
+  resGpu.setDevice(Device::CPU);
+
+  const auto expectedDims = resCpu.getDims().toVector();
+  ASSERT_EQ(resGpu.getDims().toVector(), expectedDims);
+
+  for(auto i = 0; i < resCpu.getDims().get(0); i++) {
+    for(auto j = 0; j < resCpu.getDims().get(1); j++) {
+      ASSERT_NEAR(resCpu.get(i, j), resGpu.get(i, j), 1e-4)     
+        << "Mismatch at (" << i << ", " << j << ")"
+        << " cpu=" << resCpu.get(i, j) 
+        << " gpu=" << resGpu.get(i, j);
+    }
+  }
+}
+
+TEST(CudaTensorOpsTest, MatMulThrowsWhenDimensionsNotMatched) {
+  auto t1 = TensorFunctions::Ones({2, 2});
+  auto t2 = TensorFunctions::Ones({3, 2});
+
+  ASSERT_THROW(t1.matmul(t2), std::runtime_error);
+}
+
+TEST(CudaAutogradTest, MatMul) {
+  constexpr int dimsize = 30;
+
+  // init tensors
+  auto t1 = std::make_shared<Tensor>(
+    TensorFunctions::Gaussian(
+      {10, dimsize}, 2.0, Device::CPU, true));
+
+  auto t2 = std::make_shared<Tensor>(
+    TensorFunctions::Gaussian(
+      {dimsize, 10}, 2.0, Device::CPU, true));
+
+  auto t1Gpu = std::make_shared<Tensor>(t1->createDeepCopy());
+  auto t2Gpu = std::make_shared<Tensor>(t2->createDeepCopy());
+  t1Gpu->setDevice(Device::CUDA);
+  t2Gpu->setDevice(Device::CUDA);
+
+  {
+    // compute and take loss
+    auto resCpu = cgraph::matmul(t1, t2);
+    auto lossCpu = TensorFunctions::makeSharedTensor(
+      {1}, {0.0}, Device::CPU, true);
+    for(size_t i = 0; i < resCpu->getSize(); ++i) {
+      lossCpu = cgraph::add(lossCpu, cgraph::get(resCpu, i));
+    }
+    lossCpu->backward();
+  }
+
+  {
+    auto resGpu = cgraph::matmul(t1Gpu, t2Gpu);
+    auto lossGpu = TensorFunctions::makeSharedTensor(
+      {1}, {0.0}, Device::CUDA, true);
+    for(size_t i = 0; i < resGpu->getSize(); ++i) {
+      lossGpu = cgraph::add(lossGpu, cgraph::get(resGpu, i));
+    }
+    lossGpu->backward();
+  }
+
+  // get grads and compare
+  auto t1Grads = t1->getGrads();
+  auto t2Grads = t2->getGrads();
+
+  auto t1GpuGrads = t1Gpu->getGrads();
+  auto t2GpuGrads = t2Gpu->getGrads();
+
+  for(auto i = 0; i < t1Grads->getDims().get(0); i++) {
+    for(auto j = 0; j < t1Grads->getDims().get(1); j++) {    
+      EXPECT_NEAR(t1Grads->get(i, j), t1GpuGrads->get(i, j), 1e-5)
+        << "Mismatch at (" << i << ", " << j << ")"
+        << " cpu=" << t1Grads->get(i, j) 
+        << " gpu=" << t1GpuGrads->get(i, j);
+    }
+  }
+
+  for(auto i = 0; i < t2Grads->getDims().get(0); i++) {
+    for(auto j = 0; j < t2Grads->getDims().get(1); j++) {    
+      EXPECT_NEAR(t2Grads->get(i, j), t2GpuGrads->get(i, j), 1e-5)
+        << "Mismatch at (" << i << ", " << j << ")"
+        << " cpu=" << t2Grads->get(i, j) 
+        << " gpu=" << t2GpuGrads->get(i, j);
+    }
+  }
+}
+
+TEST(CudaTensorOpsTest, MatrixTranspose1) {
+  auto t = TensorFunctions::Gaussian({200, 200}, 1.0, Device::CUDA);
+
+  auto transposed = t.createDeepCopy();
+  transposed = transposed.transpose(-1, -2);
+  transposed.setDevice(Device::CPU);
+
+  ASSERT_EQ(t.getDims().get(-1), transposed.getDims().get(-2));
+  ASSERT_EQ(t.getDims().get(-2), transposed.getDims().get(-1));
+  ASSERT_EQ(t.getDims().nDims(), transposed.getDims().nDims());
+  
+  for(auto row=0; row<t.getDims().get(-2); row++) {
+    for(auto col=0; col<t.getDims().get(-1); col++) {
+      ASSERT_NEAR(t.get(row, col), transposed.get(col, row), 1e-5)
+        << "Mismatch at (" << row << ", " << col << ")"
+        << " t=" << t.get(row, col) 
+        << " transposed=" << transposed.get(col, row);
+    }
+  }
+}
+
+// Swap first two dimensions.
+TEST(CudaTensorOpsTest, MatrixTranspose2) {
+  auto t = TensorFunctions::Gaussian({10, 20, 200}, 1.0, Device::CUDA);
+  auto transposed = t.createDeepCopy().transpose(0, 1);
+  transposed.setDevice(Device::CPU);
+
+  ASSERT_EQ(t.getDims().get(0), transposed.getDims().get(1));
+  ASSERT_EQ(t.getDims().get(1), transposed.getDims().get(0));
+  ASSERT_EQ(t.getDims().get(-1), transposed.getDims().get(-1));
+  ASSERT_EQ(t.getDims().nDims(), transposed.getDims().nDims());
+  
+  for(auto dim1=0; dim1<t.getDims().get(0); dim1++) {
+    for(auto dim2=0; dim2<t.getDims().get(1); dim2++) {
+      for(auto dim3=0; dim3<t.getDims().get(-1); dim3++) {
+        // we transposed dim1 and dim3
+        ASSERT_NEAR(t.get(dim1, dim2, dim3), transposed.get(dim2, dim1, dim3), 1e-5)
+          << "Mismatch at (" << dim1 << ", " << dim2 << ", " << dim3 << ")"
+          << " t=" << t.get(dim1, dim2, dim3)
+          << " transposed=" << transposed.get(dim2, dim1, dim3);
+      }
+    }
+  }
+}
+
+// Swap first and last dimension.
+TEST(CudaTensorOpsTest, MatrixTranspose3) {
+  auto t = TensorFunctions::Gaussian({10, 20, 200}, 1.0, Device::CUDA);
+  auto transposed = t.createDeepCopy().transpose(0, -1);
+  transposed.setDevice(Device::CPU);
+
+  ASSERT_EQ(t.getDims().get(0), transposed.getDims().get(-1));
+  ASSERT_EQ(t.getDims().get(-1), transposed.getDims().get(0));
+  ASSERT_EQ(t.getDims().get(1), transposed.getDims().get(1));
+  ASSERT_EQ(t.getDims().nDims(), transposed.getDims().nDims());
+  
+  for(auto dim1=0; dim1<t.getDims().get(0); dim1++) {
+    for(auto dim2=0; dim2<t.getDims().get(1); dim2++) {
+      for(auto dim3=0; dim3<t.getDims().get(-1); dim3++) {
+        // we transposed dim1 and dim3
+        ASSERT_NEAR(t.get(dim1, dim2, dim3), transposed.get(dim3, dim2, dim1), 1e-5)
+          << "Mismatch at (" << dim1 << ", " << dim2 << ", " << dim3 << ")"
+          << " t=" << t.get(dim1, dim2, dim3)
+          << " transposed=" << transposed.get(dim3, dim2, dim1);
+      }
+    }
+  }
+}
\ No newline at end of file
diff --git a/tests/backend/main.cpp b/tests/backend/main.cpp
new file mode 100644
index 0000000..c862d79
--- /dev/null
+++ b/tests/backend/main.cpp
@@ -0,0 +1,20 @@
+/**
+ * @file main.cpp
+ * @author Robert Baumgartner (r.baumgartner-1@tudelft.nl)
+ * @brief 
+ * @version 0.1
+ * @date 2026-05-20
+ * 
+ * @copyright Copyright (c) 2026
+ * 
+ */
+
+#include <gtest/gtest.h>
+
+#include "system/sys_functions.h"
+
+int main(int argc, char** argv) {
+  testing::InitGoogleTest(&argc, argv);
+  sys::setRandomSeed(42);
+  return RUN_ALL_TESTS();
+}
\ No newline at end of file
diff --git a/tests/backend/net_factories.h b/tests/backend/net_factories.h
new file mode 100644
index 0000000..cc6500c
--- /dev/null
+++ b/tests/backend/net_factories.h
@@ -0,0 +1,65 @@
+/**
+ * @file net_factories.h
+ * @author Robert Baumgartner (r.baumgartner-1@tudelft.nl)
+ * @brief
+ * @version 0.1
+ * @date 2026-05-20
+ *
+ * @copyright Copyright (c) 2026
+ *
+ */
+
+#pragma once
+
+#include "data_modeling/device.h"
+#include "module/networks/sequential.h"
+#include "module/layers/ff_layer.h"
+
+#include "module/activation_functions/sigmoid.h"
+#include "module/activation_functions/relu.h"
+#include "module/activation_functions/leaky_relu.h"
+#include "module/activation_functions/softmax.h"
+
+#include <memory>
+
+static std::shared_ptr<module::Sequential> makeBinaryNet(Device d = Device::CPU) {
+  auto net = std::make_shared<module::Sequential>();
+
+  net->append(std::make_shared<module::FfLayer>(2, 4, d, true, true));
+  net->append(std::make_shared<module::LeakyReLu>(0.01));
+  net->append(std::make_shared<module::FfLayer>(4, 1, d, true, true));
+  net->append(std::make_shared<module::Sigmoid>());
+
+  return net;
+}
+
+static std::shared_ptr<module::Sequential> makeBinaryNet2(Device d = Device::CPU) {
+  auto net = std::make_shared<module::Sequential>();
+
+  net->append(std::make_shared<module::FfLayer>(2, 4, d, true, true));
+  net->append(std::make_shared<module::LeakyReLu>(0.01));
+  net->append(std::make_shared<module::FfLayer>(4, 1, d, true, true));
+
+  return net;
+}
+
+static std::shared_ptr<module::Sequential> makeMulticlassNet(Device d = Device::CPU) {
+  auto net = std::make_shared<module::Sequential>();
+
+  net->append(std::make_shared<module::FfLayer>(2, 8, d, true, true));
+  net->append(std::make_shared<module::LeakyReLu>(0.01));
+  net->append(std::make_shared<module::FfLayer>(8, 3, d, true, true));
+  net->append(std::make_shared<module::Softmax>());
+
+  return net;
+}
+
+static std::shared_ptr<module::Sequential> makeMulticlassNet2(Device d = Device::CPU) {
+  auto net = std::make_shared<module::Sequential>();
+
+  net->append(std::make_shared<module::FfLayer>(2, 8, d, true, true));
+  net->append(std::make_shared<module::LeakyReLu>(0.01));
+  net->append(std::make_shared<module::FfLayer>(8, 3, d, true, true));
+
+  return net;
+}
diff --git a/tests/backend/test_computational_graph.cpp b/tests/backend/test_computational_graph.cpp
index ef07e65..489b944 100644
--- a/tests/backend/test_computational_graph.cpp
+++ b/tests/backend/test_computational_graph.cpp
@@ -1,9 +1,9 @@
 /**
- * @file test_computational_graph.cpp
+ * @file test_training.cpp
  * @author Robert Baumgartner (r.baumgartner-1@tudelft.nl)
  * @brief 
  * @version 0.1
- * @date 2026-02-16
+ * @date 2026-03-14
  * 
  * @copyright Copyright (c) 2026
  * 
@@ -11,140 +11,247 @@
 
 #include <gtest/gtest.h>
 
-#include "data_modeling/tensor.h"
 #include "data_modeling/tensor_functions.h"
-
 #include "computational_graph/tensor_ops/graph_creation.h"
 
-#include <stdexcept>
+#include "net_factories.h"
+
+#include "training/optimizers/sgd.h"
+#include "training/optimizers/rmsprop.h"
+
+#include "training/loss_functions/bce_loss.h"
+#include "training/loss_functions/crossentropy_loss.h"
+#include "training/loss_functions/bce_sigmoid_loss.h"
+#include "training/loss_functions/crossentropy_softmax_loss.h"
+
+#include "training/trainers/base_train_loop.h"
+
+using namespace std;
 
 TEST(AutogradTest, ThrowsIfNoGradientSet) {
-    auto t1 = TensorFunctions::makeSharedTensor({1}, {3.0}, false);
-    auto t2 = TensorFunctions::makeSharedTensor({1}, {2.0}, false);
+  auto t1 = TensorFunctions::makeSharedTensor({1}, {3.0}, false);
+  auto t2 = TensorFunctions::makeSharedTensor({1}, {2.0}, false);
 
-    auto loss = cgraph::add(t1, t2);
-    
-    EXPECT_THROW(loss->backward(), std::runtime_error);
+  auto loss = cgraph::add(t1, t2);
+  
+  ASSERT_THROW(loss->backward(), std::runtime_error);
 }
 
-TEST(AutogradTest, SimpleAddition) {
-    auto t1 = TensorFunctions::makeSharedTensor({1}, {3.0}, true);
-    auto t2 = TensorFunctions::makeSharedTensor({1}, {2.0}, true);
-
-    auto t3 = cgraph::add(t1, t2);
-    auto loss = cgraph::mul(t3, t3);
-    
-    loss->backward();
-    
-    EXPECT_NEAR(t1->getGrads()->get(0), 10.0, 1e-5);
-    EXPECT_NEAR(t2->getGrads()->get(0), 10.0, 1e-5);
+TEST(AutogradTest, ChainRule) {
+  auto x = TensorFunctions::makeSharedTensor({1}, {2.0}, true);
+  
+  auto y = cgraph::mul(x, x); // y = x^2
+  auto z = cgraph::add(x, y); // z = x^2 + x
+  auto loss = cgraph::mul(z, z);   // loss = (x^2 + x)^2
+  
+  loss->backward();
+  
+  // dloss/dx = 2(x^2 + x) * (2x + 1)
+  // At x=2: 2(4 + 2) * (4 + 1) = 2 * 6 * 5 = 60
+  ASSERT_NEAR(x->getGrads()->get(0), 60.0, 1e-5);
 }
 
-TEST(AutogradTest, BroadcastAdd) {
-    // gradient of broadcast add w.r.t. bias should be sum over batch dimension
-    // upstream grad: (2,3) of ones → bias grad should be (3) of twos
-    auto t1 = TensorFunctions::makeSharedTensor({2, 3}, 
-        {1.0, 2.0, 3.0,
-         4.0, 5.0, 6.0}, true);
-    auto bias = TensorFunctions::makeSharedTensor({3}, 
-        {0.0, 0.0, 0.0}, true);
-
-    auto res = cgraph::add(t1, bias);
-
-    // set upstream grad to ones and backprop
-    auto upstreamGrad = TensorFunctions::makeSharedTensor({2, 3},
-        {1.0, 1.0, 1.0,
-         1.0, 1.0, 1.0}, false);
-    res->backward();
-
-    // bias grad should be sum over batch: [2, 2, 2]
-    auto biasGrad = bias->getGrads();
-    ASSERT_DOUBLE_EQ((*biasGrad)[0], 2.0);
-    ASSERT_DOUBLE_EQ((*biasGrad)[1], 2.0);
-    ASSERT_DOUBLE_EQ((*biasGrad)[2], 2.0);
-
-    // t1 grad should be ones (add is identity for non-broadcast operand)
-    auto t1Grad = t1->getGrads();
-    for(int i = 0; i < 6; i++) {
-        ASSERT_DOUBLE_EQ((*t1Grad)[i], 1.0);
-    }
+TEST(AutogradTest, MultiVariateChainRule) {
+  auto x = TensorFunctions::makeSharedTensor({2}, {1.0, 2.0}, true);
+  
+  auto y = cgraph::mul(x, 3.0); // y = [3, 6]
+  auto loss = TensorFunctions::makeSharedTensor({1}, {0.0}, true);
+  for(int i = 0; i < y->getSize(); i++){
+    loss = cgraph::add(loss, cgraph::get(y, i));
+  }    // loss = 9
+  
+  loss->backward();
+  
+  // dloss/dx = scalar = 3
+  ASSERT_NEAR(x->getGrads()->get(0), 3.0, 1e-5);
+  ASSERT_NEAR(x->getGrads()->get(1), 3.0, 1e-5);
+
+  ASSERT_NEAR(y->getGrads()->get(0), 1.0, 1e-5);
+  ASSERT_NEAR(y->getGrads()->get(1), 1.0, 1e-5);
 }
 
-TEST(AutogradTest, ScalarMultiplication) {
-    auto t1 = TensorFunctions::makeSharedTensor({1}, {2.0}, true);
-    auto t2 = TensorFunctions::makeSharedTensor({1}, {3.0}, true);
-
-    auto t3 = cgraph::mul(t1, t2);
-    auto loss = cgraph::mul(t3, t3);
-    
-    loss->backward();
-    
-    ASSERT_DOUBLE_EQ(t1->getGrads()->get(0), 36.0);
-    ASSERT_DOUBLE_EQ(t2->getGrads()->get(0), 24.0);
+TEST(OverfitTest, BceSgdOverfitsSmallDataset) {
+  // XOR-like: 4 samples, 2 features, binary labels
+  auto x = TensorFunctions::makeSharedTensor(
+    {4, 2}, {0.0, 0.0,
+             0.0, 1.0,
+             1.0, 0.0,
+             1.0, 1.0}, false);
+
+  auto y = TensorFunctions::makeSharedTensor(
+    {4, 1}, {0.0,
+             1.0,
+             1.0,
+             0.0}, false);
+
+  auto net = makeBinaryNet();    
+  auto loss = make_shared<train::BceLoss>();
+  auto optim = make_shared<train::SgdOptimizer>(
+    net->parameters(), /*lr=*/0.05);
+
+  auto trainLoop = train::BaseTrainLoop(
+    net, loss, optim, /*epochs=*/2000, /*bsize=*/static_cast<tensorDim_t>(4));
+
+  trainLoop.run(x, y, /*shuffle=*/false, /*verbose=*/false);
+
+  // forward one more time to get final loss
+  auto pred = (*net)(x);
+  auto finalLoss = (*loss)(y, pred);
+
+  EXPECT_LT((*finalLoss)[0], 0.05f)
+    << "Network failed to overfit binary dataset\n"
+    << "Final prediction: " << *pred << "\nFinal loss: " << *finalLoss;
 }
 
-TEST(AutogradTest, MatMul) {
-    auto t1 = TensorFunctions::makeSharedTensor({2, 3}, {1, 2, 3, 4, 5, 6}, true);
-    auto t2 = TensorFunctions::makeSharedTensor({3, 2}, {1, 2, 3, 4, 5, 6}, true);
-    
-    auto t3 = cgraph::matmul(t1, t2);
+TEST(OverfitTest, BceSgdOverfitsSmallDataset_OptimizedLoss) {    
+  // XOR-like: 4 samples, 2 features, binary labels
+  auto x = TensorFunctions::makeSharedTensor(
+    {4, 2}, {0.0, 0.0,
+             0.0, 1.0,
+             1.0, 0.0,
+             1.0, 1.0}, false);
 
-    auto loss = TensorFunctions::makeSharedTensor({1}, {0.0}, true);
-    for (size_t i = 0; i < t3->getSize(); ++i) {
-        loss = cgraph::add(loss, cgraph::get(t3, i));
-    }
-    
-    loss->backward();
-    
-    EXPECT_TRUE(t1->hasGrads());
-    EXPECT_TRUE(t2->hasGrads());
-
-    // dL/dt1 = dloss/dt3 @ t2^t = Ones({2, 2}) @ t2^t
-    ASSERT_DOUBLE_EQ(t1->getGrads()->get({0, 0}), 3.0);
-    ASSERT_DOUBLE_EQ(t1->getGrads()->get({0, 1}), 7.0);
-    ASSERT_DOUBLE_EQ(t1->getGrads()->get({0, 2}), 11.0);
-    ASSERT_DOUBLE_EQ(t1->getGrads()->get({1, 0}), 3.0);
-    ASSERT_DOUBLE_EQ(t1->getGrads()->get({1, 1}), 7.0);
-    ASSERT_DOUBLE_EQ(t1->getGrads()->get({1, 2}), 11.0);
-
-    // dL/dt2 = t1^t @ dloss/dt3 = t1^t @ Ones({2, 2})
-    ASSERT_DOUBLE_EQ(t2->getGrads()->get({0, 0}), 5.0);
-    ASSERT_DOUBLE_EQ(t2->getGrads()->get({0, 1}), 5.0);
-    ASSERT_DOUBLE_EQ(t2->getGrads()->get({1, 0}), 7.0);
-    ASSERT_DOUBLE_EQ(t2->getGrads()->get({1, 1}), 7.0);
-    ASSERT_DOUBLE_EQ(t2->getGrads()->get({2, 0}), 9.0);
-    ASSERT_DOUBLE_EQ(t2->getGrads()->get({2, 1}), 9.0);
+  auto y = TensorFunctions::makeSharedTensor(
+    {4, 1}, {0.0,
+             1.0,
+             1.0,
+             0.0}, false);
+
+  auto net = makeBinaryNet2();    
+  auto loss = make_shared<train::BceSigmoidLoss>();
+  auto optim = make_shared<train::SgdOptimizer>(
+    net->parameters(), /*lr=*/0.05);
+
+  auto trainLoop = train::BaseTrainLoop(
+    net, loss, optim, /*epochs=*/2000, /*bsize=*/static_cast<tensorDim_t>(4));
+
+  trainLoop.run(x, y, /*shuffle=*/false, /*verbose=*/false);
+
+  // forward one more time to get final loss
+  auto pred = (*net)(x);
+  auto finalLoss = (*loss)(y, pred);
+
+  auto sigmoid = module::Sigmoid();
+  EXPECT_LT((*finalLoss)[0], 0.05f)
+    << "Network failed to overfit binary dataset\n"
+    << "Final prediction: " << sigmoid(*pred) << "\nFinal loss: " << *finalLoss;
 }
 
-TEST(AutogradTest, ChainRule) {
-    auto x = TensorFunctions::makeSharedTensor({1}, {2.0}, true);
-    
-    auto y = cgraph::mul(x, x); // y = x^2
-    auto z = cgraph::add(x, y); // z = x^2 + x
-    auto loss = cgraph::mul(z, z);   // loss = (x^2 + x)^2
-    
-    loss->backward();
-    
-    // dloss/dx = 2(x^2 + x) * (2x + 1)
-    // At x=2: 2(4 + 2) * (4 + 1) = 2 * 6 * 5 = 60
-    ASSERT_DOUBLE_EQ(x->getGrads()->get(0), 60.0);
+TEST(OverfitTest, CrossEntropyRMSPropOverfitsSmallDataset) {
+  // 6 samples, 2 features, 3 classes
+  auto x = TensorFunctions::makeSharedTensor(
+    {6, 2}, {1.0, 0.0,
+             1.0, 0.1,
+             0.0, 1.0,
+             0.1, 1.0,
+             0.5, 0.5,
+             0.4, 0.6}, false);
+
+  // one-hot encoded labels
+  auto y = TensorFunctions::makeSharedTensor(
+    {6, 3}, {1.0, 0.0, 0.0,
+             1.0, 0.0, 0.0,
+             0.0, 1.0, 0.0,
+             0.0, 1.0, 0.0,
+             0.0, 0.0, 1.0,
+             0.0, 0.0, 1.0}, false);
+
+  auto net = makeMulticlassNet();
+  auto loss = make_shared<train::CrossEntropyLoss>();
+  auto optim = make_shared<train::RmsPropOptimizer>(
+    net->parameters(), /*lr=*/0.001, /*decay=*/0.95);
+
+  auto trainLoop = train::BaseTrainLoop(
+    net, loss, optim, /*epochs=*/2000, /*bsize=*/6);
+
+  trainLoop.run(x, y, /*shuffle=*/false, /*verbose=*/false);
+
+  auto pred = (*net)(x);
+  auto finalLoss = (*loss)(y, pred);
+
+  EXPECT_LT((*finalLoss)[0], 0.05f)
+    << "Network failed to overfit multiclass dataset"
+    << "Final prediction: " << *pred << "\nFinal loss: " << *finalLoss;
 }
 
-TEST(AutogradTest, MultiVariateChainRule) {
-    auto x = TensorFunctions::makeSharedTensor({2}, {1.0, 2.0}, true);
-    
-    auto y = cgraph::mul(x, 3.0); // y = [3, 6]
-    auto loss = TensorFunctions::makeSharedTensor({1}, {0.0}, true);
-    for(int i=0; i<y->getSize(); i++){
-        loss = cgraph::add(loss, cgraph::get(y, i));
-    }    // loss = 9
-    
-    loss->backward();
-    
-    // dloss/dx = scalar = 3
-    ASSERT_DOUBLE_EQ(x->getGrads()->get(0), 3.0);
-    ASSERT_DOUBLE_EQ(x->getGrads()->get(1), 3.0);
-
-    ASSERT_DOUBLE_EQ(y->getGrads()->get(0), 1.0);
-    ASSERT_DOUBLE_EQ(y->getGrads()->get(1), 1.0);
+TEST(OverfitTest, CrossEntropyRMSPropOverfitsSmallDataset_OptimizedLoss) {
+  // 6 samples, 2 features, 3 classes
+  auto x = TensorFunctions::makeSharedTensor(
+    {6, 2}, {1.0, 0.0,
+             1.0, 0.1,
+             0.0, 1.0,
+             0.1, 1.0,
+             0.5, 0.5,
+             0.4, 0.6}, false);
+
+  // one-hot encoded labels
+  auto y = TensorFunctions::makeSharedTensor(
+    {6, 3}, {1.0, 0.0, 0.0,
+             1.0, 0.0, 0.0,
+             0.0, 1.0, 0.0,
+             0.0, 1.0, 0.0,
+             0.0, 0.0, 1.0,
+             0.0, 0.0, 1.0}, false);
+
+  auto net = makeMulticlassNet2();
+  auto loss = make_shared<train::CrossEntropySoftmaxLoss>();
+  auto optim = make_shared<train::RmsPropOptimizer>(
+    net->parameters(), /*lr=*/0.001, /*decay=*/0.95);
+
+  auto trainLoop = train::BaseTrainLoop(
+    net, loss, optim, /*epochs=*/2000, /*bsize=*/6);
+
+  trainLoop.run(x, y, /*shuffle=*/false, /*verbose=*/false);
+
+  auto pred = (*net)(x);
+  auto finalLoss = (*loss)(y, pred);
+
+  auto softmax = module::Softmax();
+  EXPECT_LT((*finalLoss)[0], 0.05f)
+    << "Network failed to overfit multiclass dataset"
+    << "Final prediction: " << softmax(*pred) << "\nFinal loss: " << *finalLoss;
+}
+
+TEST(OptimizerTest, ZeroGrad_ClearsAllGradients) {
+  auto x = TensorFunctions::makeSharedTensor(
+    {4, 2}, {0.0, 0.0, 0.0, 1.0, 1.0, 0.0, 1.0, 1.0}, false);
+  auto y = TensorFunctions::makeSharedTensor(
+    {4, 1}, {0.0, 1.0, 1.0, 0.0}, false);
+
+  auto net = makeBinaryNet();
+  auto loss = std::make_shared<train::BceSigmoidLoss>();
+  auto optim = std::make_shared<train::SgdOptimizer>(
+    net->parameters(), 0.01f);
+
+  // one forward/backward pass to populate gradients
+  auto pred = (*net)(x);
+  auto l = (*loss)(y, pred);
+  l->backward();
+
+  // verify gradients are non-zero before zeroing
+  bool anyNonZero = false;
+  for(auto& p : net->parameters()) {
+    if(p->getGrads()) {
+      for(tensorSize_t i = 0; i < p->getGrads()->getSize(); i++) {
+        if((*p->getGrads())[i] != 0.0f) {
+          anyNonZero = true;
+          break;
+        }
+      }
+    }
+  }
+  EXPECT_TRUE(anyNonZero) << "Expected some non-zero gradients before zeroGrad";
+
+  // zero gradients
+  optim->zeroGrad();
+
+  // verify all gradients are zero after zeroing
+  for(auto& p : net->parameters()) {
+    if(p->getGrads()) {
+      for(tensorSize_t i = 0; i < p->getGrads()->getSize(); i++) {
+        ASSERT_NEAR((*p->getGrads())[i], 0.0f, 1e-5)
+          << "Gradient not zeroed at index " << i;
+      }
+    }
+  }
 }
\ No newline at end of file
diff --git a/tests/backend/test_data_modeling.cpp b/tests/backend/test_data_modeling.cpp
deleted file mode 100644
index 69f585b..0000000
--- a/tests/backend/test_data_modeling.cpp
+++ /dev/null
@@ -1,303 +0,0 @@
-/**
- * @file test_data_modeling.cpp
- * @author Robert Baumgartner (r.baumgartner-1@tudelft.nl)
- * @brief 
- * @version 0.1
- * @date 2026-02-16
- * 
- * @copyright Copyright (c) 2026
- * 
- */
-
-#include <gtest/gtest.h>
-
-#include "data_modeling/tensor.h"
-#include "data_modeling/tensor_functions.h"
-
-#include <stdexcept>
-
-TEST(TensorOpsTest, TestCtor) {
-  auto t = Tensor({2, 2}, {2.0, 3.0, 4.0, 5.0}, Device::CPU, false);
-
-  ASSERT_EQ(t.getDims(), Dimension({2, 2}));
-  ASSERT_EQ(t.getDevice(), Device::CPU);
-  ASSERT_TRUE(!t.getRequiresGrad());
-
-  ASSERT_DOUBLE_EQ(t.get(0, 0), 2.0);
-  ASSERT_DOUBLE_EQ(t.get(0, 1), 3.0);
-  ASSERT_DOUBLE_EQ(t.get(1, 0), 4.0);
-  ASSERT_DOUBLE_EQ(t.get(1, 1), 5.0);
-}
-
-TEST(TensorOpsTest, ScalarAddWorks) {
-  auto t1 = TensorFunctions::Ones({2, 2}, false);
-
-  auto res = t1 + 1.5;
-
-  constexpr ftype sum = 2.5;
-  for(auto i=0; i<t1.getDims().get(0); i++) {
-    for(auto j=0; j<t1.getDims().get(1); j++) {
-      ASSERT_DOUBLE_EQ(res.get(i, j), sum);
-    }
-  }
-}
-
-TEST(TensorOpsTest, TensorAddWorks) {
-  auto t1 = TensorFunctions::Ones({2, 2}, false);
-  auto t2 = TensorFunctions::Ones({2, 2}, false) * 4;
-
-  auto res = t1 + t2;
-
-  constexpr ftype sum = 5.0;
-  for(auto i=0; i<t1.getDims().get(0); i++) {
-    for(auto j=0; j<t1.getDims().get(1); j++) {
-      ASSERT_DOUBLE_EQ(res.get(i, j), sum);
-    }
-  }
-}
-
-TEST(TensorOpsTest, TensorAddCanBroadCast) {
-  auto t1 = TensorFunctions::Ones({3, 2, 2}, false);
-  auto t2 = Tensor({2}, {2, 3}, false);
-
-  auto res = t1 + t2;
-
-  ASSERT_EQ(res.getDims(), t1.getDims());
-  
-  for(auto i=0; i<res.getDims().get(0); i++) {
-    for(auto j=0; j<res.getDims().get(1); j++) {
-      ASSERT_DOUBLE_EQ(res.get(i, j, 0), 3.0);
-      ASSERT_DOUBLE_EQ(res.get(i, j, 1), 4.0);
-    }
-  }
-}
-
-TEST(TensorOpsTest, BroadcastAdd_2D) {
-    // (2,3) + (3) 
-    auto t1 = Tensor({2, 3}, {1.0, 2.0, 3.0,
-                               4.0, 5.0, 6.0}, false);
-    auto t2 = Tensor({3}, {10.0, 20.0, 30.0}, false);
-    auto res = t1 + t2;
-
-    // expected: each row of t1 gets t2 added elementwise
-    ASSERT_DOUBLE_EQ(res.get(0, 0), 11.0);
-    ASSERT_DOUBLE_EQ(res.get(0, 1), 22.0);
-    ASSERT_DOUBLE_EQ(res.get(0, 2), 33.0);
-    ASSERT_DOUBLE_EQ(res.get(1, 0), 14.0);
-    ASSERT_DOUBLE_EQ(res.get(1, 1), 25.0);
-    ASSERT_DOUBLE_EQ(res.get(1, 2), 36.0);
-}
-
-TEST(TensorOpsTest, TensorAddBroadcastNotCommutative) {
-  auto t1 = TensorFunctions::Ones({3, 2, 2}, false);
-  auto t2 = Tensor({2}, {2, 3}, false);
-
-  EXPECT_THROW(t2 + t1, std::invalid_argument);
-}
-
-TEST(TensorOpsTest, TensorAddThrowsOnDimMismatch) {
-  auto t1 = TensorFunctions::Ones({2, 2}, false);
-  auto t2 = TensorFunctions::Ones({2, 3}, false) * 4;
-
-  EXPECT_THROW(t1 + t2, std::invalid_argument);
-}
-
-TEST(TensorOpsTest, ScalarMulWorks) {
-  auto t1 = TensorFunctions::Ones({2, 2}, false);
-
-  constexpr ftype f = 2.5;
-  auto res = t1 * f;
-    
-  for(auto i=0; i<t1.getDims().get(0); i++) {
-    for(auto j=0; j<t1.getDims().get(1); j++) {
-      ASSERT_DOUBLE_EQ(res.get(i, j), f);
-    }
-  }
-}
-
-TEST(TensorOpsTest, MatrixAddGivesCorrectResults) {
-  auto t1 = TensorFunctions::Ones({2, 2}, false);
-  auto t2 = TensorFunctions::Ones({2, 2}, false);
-    
-  auto res = t1 + t2;
-    
-  constexpr ftype resSum = 2.0;
-
-  for(auto i=0; i<t1.getDims().get(0); i++) {
-    for(auto j=0; j<t1.getDims().get(1); j++) {
-      ASSERT_DOUBLE_EQ(res.get(i, j), resSum);
-    }
-  }
-}
-
-TEST(TensorOpsTest, ElementwiseMulGivesCorrectResults) {
-  constexpr ftype factor = 0.5;
-  auto t1 = TensorFunctions::Ones({2, 2}, false);
-  auto t2 = TensorFunctions::Ones({2, 2}, false) * 0.5;
-    
-  auto res = t1 * t2;
-    
-  for(auto i=0; i<t1.getDims().get(0); i++) {
-    for(auto j=0; j<t1.getDims().get(1); j++) {
-      ASSERT_DOUBLE_EQ(res.get(i, j), factor);
-    }
-  }
-}
-
-TEST(TensorOpsTest, ElementwiseMulThrowsOnDimensionMismatch) {
-  constexpr ftype factor = 0.5;
-  auto t1 = TensorFunctions::Ones({2, 2}, false);
-  auto t2 = TensorFunctions::Ones({2, 3}, false) * 0.5;
-    
-  EXPECT_THROW(t1 * t2, std::invalid_argument);
-}
-
-TEST(TensorOpsTest, MatMulGivesCorrectValues1) {
-  auto t1 = TensorFunctions::Ones({3, 2}, false);
-  auto t2 = TensorFunctions::Ones({2, 6}, false) * 1.5;
-    
-  auto res = t1.matmul(t2);
-  auto expectedDims = std::vector<tensorDim_t>{3, 6};
-  ASSERT_EQ(res.getDims().toVector(), expectedDims);
-
-  constexpr ftype resSum = 3.0;
-  for(auto i=0; i<t1.getDims().get(0); i++) {
-    for(auto j=0; j<t1.getDims().get(1); j++) {
-      ASSERT_DOUBLE_EQ(res.get(i, j), resSum);
-    }
-  }
-}
-
-TEST(TensorOpsTest, MatMulGivesCorrectValues2) {
-  auto t1 = Tensor({2, 2}, false);
-  auto t2 = Tensor({2, 2}, false);
-
-  auto cmpRes = Tensor({2, 2}, false);
-
-  auto populateTensor = [](Tensor& t, ftype v1, ftype v2, ftype v3, ftype v4) {
-    t.set(v1, {0, 0});
-    t.set(v2, {0, 1});
-    t.set(v3, {1, 0});
-    t.set(v4, {1, 1});
-  };
-
-  populateTensor(t1, 1, 2, 3, 4);
-  populateTensor(t2, 5, 6, 7, 8);
-  populateTensor(cmpRes, 19, 22, 43, 50);
-
-  auto res = t1.matmul(t2);
-
-  auto expectedDims = std::vector<tensorDim_t>{2, 2};
-  ASSERT_EQ(res.getDims().toVector(), expectedDims);
-
-  constexpr ftype resSum = 3.0;
-  for(auto i=0; i<t1.getDims().get(0); i++) {
-    for(auto j=0; j<t1.getDims().get(1); j++) {
-      ASSERT_DOUBLE_EQ(res.get(i, j), cmpRes.get(i, j));
-    }
-  }
-}
-
-TEST(TensorOpsTest, MatMulThrowsWhenDimensionsNotMatched) {
-  auto t1 = TensorFunctions::Ones({2, 2}, false);
-  auto t2 = TensorFunctions::Ones({3, 2}, false);
-
-  EXPECT_THROW(t1.matmul(t2), std::runtime_error);
-}
-
-TEST(TensorOpsTest, TransposeWorksAsIntended1) {
-  auto t = TensorFunctions::Gaussian({3, 2}, false);
-  auto transposed = t.transpose(-1, -2);
-  
-  ASSERT_EQ(t.getDims().get(-1), transposed.getDims().get(-2));
-  ASSERT_EQ(t.getDims().get(-2), transposed.getDims().get(-1));
-  ASSERT_EQ(t.getDims().nDims(), transposed.getDims().nDims());
-  
-  for(auto row=0; row<t.getDims().get(-2); row++) {
-    for(auto col=0; col<t.getDims().get(-1); col++) {
-      ASSERT_DOUBLE_EQ(t.get(row, col), transposed.get(col, row));
-    }
-  }
-}
-
-/**
- * @brief Swap first two dimensions.
- */
-TEST(TensorOpsTest, TransposeWorksAsIntended2) {
-  auto t = TensorFunctions::Gaussian({3, 2, 5}, false);
-  auto transposed = t.transpose(0, 1);
-
-  ASSERT_EQ(t.getDims().get(0), transposed.getDims().get(1));
-  ASSERT_EQ(t.getDims().get(1), transposed.getDims().get(0));
-  ASSERT_EQ(t.getDims().get(-1), transposed.getDims().get(-1));
-  ASSERT_EQ(t.getDims().nDims(), transposed.getDims().nDims());
-  
-  for(auto dim1=0; dim1<t.getDims().get(0); dim1++) {
-    for(auto dim2=0; dim2<t.getDims().get(1); dim2++) {
-      for(auto dim3=0; dim3<t.getDims().get(-1); dim3++) {
-        // we transposed dim1 and dim3
-        ASSERT_DOUBLE_EQ(t.get(dim1, dim2, dim3), transposed.get(dim2, dim1, dim3));
-      }
-    }
-  }
-}
-
-/**
- * @brief Swap first and last dimension.
- */
-TEST(TensorOpsTest, TransposeWorksAsIntended3) {
-  auto t = TensorFunctions::Gaussian({3, 2, 5}, false);
-  auto transposed = t.transpose(0, -1);
-
-  ASSERT_EQ(t.getDims().get(0), transposed.getDims().get(-1));
-  ASSERT_EQ(t.getDims().get(-1), transposed.getDims().get(0));
-  ASSERT_EQ(t.getDims().get(1), transposed.getDims().get(1));
-  ASSERT_EQ(t.getDims().nDims(), transposed.getDims().nDims());
-  
-  for(auto dim1=0; dim1<t.getDims().get(0); dim1++) {
-    for(auto dim2=0; dim2<t.getDims().get(1); dim2++) {
-      for(auto dim3=0; dim3<t.getDims().get(-1); dim3++) {
-        // we transposed dim1 and dim3
-        ASSERT_DOUBLE_EQ(t.get(dim1, dim2, dim3), transposed.get(dim3, dim2, dim1));
-      }
-    }
-  }
-}
-
-TEST(TensorOpsTest, TransposeThisWorksAsIntended1) {
-  auto t = TensorFunctions::Gaussian({3, 2}, false);
-  auto tCopy = t.createDeepCopy();
-
-  t.transposeThis();
-
-  ASSERT_EQ(t.getDims().get(-1), tCopy.getDims().get(-2));
-  ASSERT_EQ(t.getDims().get(-2), tCopy.getDims().get(-1));
-  ASSERT_EQ(t.getDims().nDims(), tCopy.getDims().nDims());
-  
-  for(auto row=0; row<t.getDims().get(-2); row++) {
-    for(auto col=0; col<t.getDims().get(-2); col++) {
-      ASSERT_DOUBLE_EQ(t.get(row, col), tCopy.get(col, row));
-    }
-  }
-}
-
-TEST(TensorOpsTest, TransposeThisWorksAsIntended2) {
-  auto t = TensorFunctions::Gaussian({3, 2, 5}, false);
-  auto tCopy = t.createDeepCopy();
-  
-  t.transposeThis(0, -1);
-
-  ASSERT_EQ(t.getDims().get(0), tCopy.getDims().get(-1));
-  ASSERT_EQ(t.getDims().get(-1), tCopy.getDims().get(0));
-  ASSERT_EQ(t.getDims().get(1), tCopy.getDims().get(1));
-  ASSERT_EQ(t.getDims().nDims(), tCopy.getDims().nDims());
-  
-  for(auto dim1=0; dim1<t.getDims().get(0); dim1++) {
-    for(auto dim2=0; dim2<t.getDims().get(1); dim2++) {
-      for(auto dim3=0; dim3<t.getDims().get(-1); dim3++) {
-        // we transposed dim1 and dim3
-        ASSERT_DOUBLE_EQ(t.get(dim1, dim2, dim3), tCopy.get(dim3, dim2, dim1));
-      }
-    }
-  }
-}
\ No newline at end of file
diff --git a/tests/backend/test_losses.cpp b/tests/backend/test_losses.cpp
index 180d6ae..efde0ef 100644
--- a/tests/backend/test_losses.cpp
+++ b/tests/backend/test_losses.cpp
@@ -21,228 +21,224 @@
 
 using namespace train;
 
-static constexpr ftype kTol = 1e-4f;
-
-// ─── CrossEntropy ────────────────────────────────────────────────────────────
+static constexpr ftype delta = 1e-4f;
 
 TEST(LossTest, CrossEntropyFoward) {
-    auto y = TensorFunctions::makeSharedTensor(
-        {2, 3}, {1.0, 0.0, 0.0,
-                 0.0, 1.0, 0.0}, false);
+  auto y = TensorFunctions::makeSharedTensor(
+    {2, 3}, {1.0, 0.0, 0.0,
+        0.0, 1.0, 0.0}, false);
 
-    auto ypred = TensorFunctions::makeSharedTensor(
-        {2, 3}, {0.7, 0.2, 0.1,
-                 0.1, 0.8, 0.1}, true);
+  auto ypred = TensorFunctions::makeSharedTensor(
+    {2, 3}, {0.7, 0.2, 0.1,
+        0.1, 0.8, 0.1}, true);
 
-    CrossEntropyLoss loss;
-    auto result = loss(y, ypred);
+  CrossEntropyLoss loss;
+  auto result = loss(y, ypred);
 
-    // expected: -( log(0.7) + log(0.8) ) / 2 = 0.2899
-    const ftype expected = -(std::log(0.7f) + std::log(0.8f)) / 2.0f;
-    EXPECT_NEAR((*result)[0], expected, kTol);
+  // expected: -( log(0.7) + log(0.8) ) / 2 = 0.2899
+  const ftype expected = -(std::log(0.7f) + std::log(0.8f)) / 2.0f;
+  ASSERT_NEAR((*result)[0], expected, delta);
 }
 
 TEST(LossTest, CrossEntropyPerfectPrediction) {
-    auto y = TensorFunctions::makeSharedTensor(
-        {2, 3}, {1.0, 0.0, 0.0,
-                 0.0, 1.0, 0.0}, false);
+  auto y = TensorFunctions::makeSharedTensor(
+    {2, 3}, {1.0, 0.0, 0.0,
+        0.0, 1.0, 0.0}, false);
 
-    // near-perfect predictions — can't use exactly 1.0 due to log(0)
-    auto ypred = TensorFunctions::makeSharedTensor(
-        {2, 3}, {0.999, 0.0005, 0.0005,
-                 0.0005, 0.999, 0.0005}, true);
+  // near-perfect predictions — can't use exactly 1.0 due to log(0)
+  auto ypred = TensorFunctions::makeSharedTensor(
+    {2, 3}, {0.999, 0.0005, 0.0005,
+        0.0005, 0.999, 0.0005}, true);
 
-    CrossEntropyLoss loss;
-    auto result = loss(y, ypred);
+  CrossEntropyLoss loss;
+  auto result = loss(y, ypred);
 
-    // loss should be very small
-    EXPECT_LT((*result)[0], 0.01f);
+  // loss should be very small
+  ASSERT_LT((*result)[0], 0.01f);
 }
 
 TEST(LossTest, CrossEntropyUniformPrediction) {
-    // uniform prediction should give log(3) ~ 1.0986
-    auto y = TensorFunctions::makeSharedTensor(
-        {1, 3}, {1.0, 0.0, 0.0}, false);
+  // uniform prediction should give log(3) ~ 1.0986
+  auto y = TensorFunctions::makeSharedTensor(
+    {1, 3}, {1.0, 0.0, 0.0}, false);
 
-    auto ypred = TensorFunctions::makeSharedTensor(
-        {1, 3}, {1.0f/3, 1.0f/3, 1.0f/3}, true);
+  auto ypred = TensorFunctions::makeSharedTensor(
+    {1, 3}, {1.0f/3, 1.0f/3, 1.0f/3}, true);
 
-    CrossEntropyLoss loss;
-    auto result = loss(y, ypred);
+  CrossEntropyLoss loss;
+  auto result = loss(y, ypred);
 
-    EXPECT_NEAR((*result)[0], std::log(3.0f), kTol);
+  ASSERT_NEAR((*result)[0], std::log(3.0f), delta);
 }
 
 TEST(LossTest, CrossEntropyThrowsOnDimMismatch) {
-    auto y = TensorFunctions::makeSharedTensor(
-        {2, 3}, {1.0, 0.0, 0.0, 0.0, 1.0, 0.0}, false);
-    auto ypred = TensorFunctions::makeSharedTensor(
-        {2, 2}, {0.5, 0.5, 0.5, 0.5}, true);
+  auto y = TensorFunctions::makeSharedTensor(
+    {2, 3}, {1.0, 0.0, 0.0, 0.0, 1.0, 0.0}, false);
+  auto ypred = TensorFunctions::makeSharedTensor(
+    {2, 2}, {0.5, 0.5, 0.5, 0.5}, true);
 
-    CrossEntropyLoss loss;
-    EXPECT_THROW(loss(y, ypred), std::invalid_argument);
+  CrossEntropyLoss loss;
+  ASSERT_THROW(loss(y, ypred), std::invalid_argument);
 }
 
 TEST(LossTest, CrossEntropyBackward) {
-    // y = [[1,0,0],[0,1,0]], ypred = [[0.7,0.2,0.1],[0.1,0.8,0.1]]
-    // grad CE w.r.t. ypred[b,i] = -y[b,i] / (ypred[b,i] * n)
-    // grad[0,0] = -1/(0.7*2) = -0.7143
-    // grad[0,1] =  0
-    // grad[0,2] =  0
-    // grad[1,0] =  0
-    // grad[1,1] = -1/(0.8*2) = -0.625
-    // grad[1,2] =  0
-    auto y = TensorFunctions::makeSharedTensor(
-        {2, 3}, {1.0, 0.0, 0.0,
-                 0.0, 1.0, 0.0}, false);
-    auto ypred = TensorFunctions::makeSharedTensor(
-        {2, 3}, {0.7, 0.2, 0.1,
-                 0.1, 0.8, 0.1}, true);
-
-    CrossEntropyLoss loss;
-    auto result = loss(y, ypred);
-    result->backward();
-
-    auto grads = ypred->getGrads();
-    EXPECT_NEAR((*grads)[0], -0.7143f, kTol);
-    EXPECT_NEAR((*grads)[1],  0.0f,    kTol);
-    EXPECT_NEAR((*grads)[2],  0.0f,    kTol);
-    EXPECT_NEAR((*grads)[3],  0.0f,    kTol);
-    EXPECT_NEAR((*grads)[4], -0.625f,  kTol);
-    EXPECT_NEAR((*grads)[5],  0.0f,    kTol);
+  // y = [[1,0,0],[0,1,0]], ypred = [[0.7,0.2,0.1],[0.1,0.8,0.1]]
+  // grad CE w.r.t. ypred[b,i] = -y[b,i] / (ypred[b,i] * n)
+  // grad[0,0] = -1/(0.7*2) = -0.7143
+  // grad[0,1] =  0
+  // grad[0,2] =  0
+  // grad[1,0] =  0
+  // grad[1,1] = -1/(0.8*2) = -0.625
+  // grad[1,2] =  0
+  auto y = TensorFunctions::makeSharedTensor(
+    {2, 3}, {1.0, 0.0, 0.0,
+        0.0, 1.0, 0.0}, false);
+  auto ypred = TensorFunctions::makeSharedTensor(
+    {2, 3}, {0.7, 0.2, 0.1,
+        0.1, 0.8, 0.1}, true);
+
+  CrossEntropyLoss loss;
+  auto result = loss(y, ypred);
+  result->backward();
+
+  auto grads = ypred->getGrads();
+  ASSERT_NEAR((*grads)[0], -0.7143f, delta);
+  ASSERT_NEAR((*grads)[1],  0.0f,    delta);
+  ASSERT_NEAR((*grads)[2],  0.0f,    delta);
+  ASSERT_NEAR((*grads)[3],  0.0f,    delta);
+  ASSERT_NEAR((*grads)[4], -0.625f,  delta);
+  ASSERT_NEAR((*grads)[5],  0.0f,    delta);
 }
 
-// ─── BCE ─────────────────────────────────────────────────────────────────────
-
 TEST(LossTest, BceForward) {
-    auto y = TensorFunctions::makeSharedTensor(
-        {4, 1}, {0.0, 1.0, 1.0, 0.0}, false);
+  auto y = TensorFunctions::makeSharedTensor(
+    {4, 1}, {0.0, 1.0, 1.0, 0.0}, false);
 
-    auto ypred = TensorFunctions::makeSharedTensor(
-        {4, 1}, {0.1, 0.9, 0.8, 0.2}, true);
+  auto ypred = TensorFunctions::makeSharedTensor(
+    {4, 1}, {0.1, 0.9, 0.8, 0.2}, true);
 
-    BceLoss loss;
-    auto result = loss(y, ypred);
+  BceLoss loss;
+  auto result = loss(y, ypred);
 
-    // expected: -( log(0.9) + log(0.9) + log(0.8) + log(0.8) ) / 4 = 0.1643
-    const ftype expected = -(std::log(0.9f) + std::log(0.9f) + 
-                              std::log(0.8f) + std::log(0.8f)) / 4.0f;
-    EXPECT_NEAR((*result)[0], expected, kTol);
+  // expected: -( log(0.9) + log(0.9) + log(0.8) + log(0.8) ) / 4 = 0.1643
+  const ftype expected = -(std::log(0.9f) + std::log(0.9f) + 
+               std::log(0.8f) + std::log(0.8f)) / 4.0f;
+  ASSERT_NEAR((*result)[0], expected, delta);
 }
 
 TEST(LossTest, BcePerfectPrediction) {
-    auto y = TensorFunctions::makeSharedTensor(
-        {2, 1}, {1.0, 0.0}, false);
+  auto y = TensorFunctions::makeSharedTensor(
+    {2, 1}, {1.0, 0.0}, false);
 
-    auto ypred = TensorFunctions::makeSharedTensor(
-        {2, 1}, {0.999, 0.001}, true);
+  auto ypred = TensorFunctions::makeSharedTensor(
+    {2, 1}, {0.999, 0.001}, true);
 
-    BceLoss loss;
-    auto result = loss(y, ypred);
+  BceLoss loss;
+  auto result = loss(y, ypred);
 
-    EXPECT_LT((*result)[0], 0.01f);
+  ASSERT_LT((*result)[0], 0.01f);
 }
 
 TEST(LossTest, BceRandomPrediction) {
-    // ypred = 0.5 for all -> loss = log(2) ~ 0.6931
-    auto y = TensorFunctions::makeSharedTensor(
-        {2, 1}, {1.0, 0.0}, false);
+  // ypred = 0.5 for all -> loss = log(2) ~ 0.6931
+  auto y = TensorFunctions::makeSharedTensor(
+    {2, 1}, {1.0, 0.0}, false);
 
-    auto ypred = TensorFunctions::makeSharedTensor(
-        {2, 1}, {0.5, 0.5}, true);
+  auto ypred = TensorFunctions::makeSharedTensor(
+    {2, 1}, {0.5, 0.5}, true);
 
-    BceLoss loss;
-    auto result = loss(y, ypred);
+  BceLoss loss;
+  auto result = loss(y, ypred);
 
-    EXPECT_NEAR((*result)[0], std::log(2.0f), kTol);
+  ASSERT_NEAR((*result)[0], std::log(2.0f), delta);
 }
 
 TEST(LossTest, BceThrowsOnDimMismatch) {
-    auto y = TensorFunctions::makeSharedTensor(
-        {2, 1}, {1.0, 0.0}, false);
-    auto ypred = TensorFunctions::makeSharedTensor(
-        {3, 1}, {0.5, 0.5, 0.5}, true);
+  auto y = TensorFunctions::makeSharedTensor(
+    {2, 1}, {1.0, 0.0}, false);
+  auto ypred = TensorFunctions::makeSharedTensor(
+    {3, 1}, {0.5, 0.5, 0.5}, true);
 
-    BceLoss loss;
-    EXPECT_THROW(loss(y, ypred), std::invalid_argument);
+  BceLoss loss;
+  ASSERT_THROW(loss(y, ypred), std::invalid_argument);
 }
 
 TEST(LossTest, BceNoInfOrNanOnNearZeroPred) {
-    auto y = TensorFunctions::makeSharedTensor(
-        {1, 1}, {1.0}, false);
-    auto ypred = TensorFunctions::makeSharedTensor(
-        {1, 1}, {0.0}, true);
+  auto y = TensorFunctions::makeSharedTensor(
+    {1, 1}, {1.0}, false);
+  auto ypred = TensorFunctions::makeSharedTensor(
+    {1, 1}, {0.0}, true);
 
-    BceLoss loss;
-    auto result = loss(y, ypred);
+  BceLoss loss;
+  auto result = loss(y, ypred);
 
-    // clipping prevents log(0)
-    EXPECT_FALSE(std::isinf((*result)[0]));
+  // clipping prevents log(0)
+  ASSERT_FALSE(std::isinf((*result)[0]));
 }
 
 TEST(LossTest, BceBackward) {
-    // y = [1, 0], ypred = [0.8, 0.3]
-    // grad BCE w.r.t. ypred_i = (-y/ypred + (1-y)/(1-ypred)) / n
-    // grad[0] = (-1/0.8 + 0) / 2 = -0.625
-    // grad[1] = (0 + 1/0.7)  / 2 =  0.7143
-    auto y = TensorFunctions::makeSharedTensor(
-        {2, 1}, {1.0, 0.0}, false);
-    auto ypred = TensorFunctions::makeSharedTensor(
-        {2, 1}, {0.8, 0.3}, true);
-
-    BceLoss loss;
-    auto result = loss(y, ypred);
-    result->backward();
-
-    auto grads = ypred->getGrads();
-    EXPECT_NEAR((*grads)[0], -0.625f,  kTol);
-    EXPECT_NEAR((*grads)[1],  0.7143f, kTol);
+  // y = [1, 0], ypred = [0.8, 0.3]
+  // grad BCE w.r.t. ypred_i = (-y/ypred + (1-y)/(1-ypred)) / n
+  // grad[0] = (-1/0.8 + 0) / 2 = -0.625
+  // grad[1] = (0 + 1/0.7)  / 2 =  0.7143
+  auto y = TensorFunctions::makeSharedTensor(
+    {2, 1}, {1.0, 0.0}, false);
+  auto ypred = TensorFunctions::makeSharedTensor(
+    {2, 1}, {0.8, 0.3}, true);
+
+  BceLoss loss;
+  auto result = loss(y, ypred);
+  result->backward();
+
+  auto grads = ypred->getGrads();
+  ASSERT_NEAR((*grads)[0], -0.625f,  delta);
+  ASSERT_NEAR((*grads)[1],  0.7143f, delta);
 }
 
 TEST(LossTest, RmseForward) {
-    // y = [1, 2, 3], ypred = [1.5, 2.5, 2.5]
-    // diffs = [-0.5, -0.5, 0.5]
-    // MSE = (0.25 + 0.25 + 0.25) / 3 = 0.25
-    // RMSE = 0.5
-    auto y = TensorFunctions::makeSharedTensor(
-        {3}, {1.0, 2.0, 3.0}, false);
-    auto ypred = TensorFunctions::makeSharedTensor(
-        {3}, {1.5, 2.5, 2.5}, true);
-
-    auto loss = RmseLoss{};
-    auto result = loss(y, ypred);
-
-    EXPECT_NEAR((*result)[0], 0.5f, kTol);
+  // y = [1, 2, 3], ypred = [1.5, 2.5, 2.5]
+  // diffs = [-0.5, -0.5, 0.5]
+  // MSE = (0.25 + 0.25 + 0.25) / 3 = 0.25
+  // RMSE = 0.5
+  auto y = TensorFunctions::makeSharedTensor(
+    {3}, {1.0, 2.0, 3.0}, false);
+  auto ypred = TensorFunctions::makeSharedTensor(
+    {3}, {1.5, 2.5, 2.5}, true);
+
+  auto loss = RmseLoss{};
+  auto result = loss(y, ypred);
+
+  ASSERT_NEAR((*result)[0], 0.5f, delta);
 }
 
 TEST(LossTest, RmsePerfectPrediction) {
-    auto y = TensorFunctions::makeSharedTensor(
-        {3}, {1.0, 2.0, 3.0}, false);
-    auto ypred = TensorFunctions::makeSharedTensor(
-        {3}, {1.0, 2.0, 3.0}, true);
+  auto y = TensorFunctions::makeSharedTensor(
+    {3}, {1.0, 2.0, 3.0}, false);
+  auto ypred = TensorFunctions::makeSharedTensor(
+    {3}, {1.0, 2.0, 3.0}, true);
 
-    RmseLoss loss;
-    auto result = loss(y, ypred);
+  RmseLoss loss;
+  auto result = loss(y, ypred);
 
-    EXPECT_NEAR((*result)[0], 0.0f, kTol);
+  ASSERT_NEAR((*result)[0], 0.0f, delta);
 }
 
 TEST(LossTest, RmseBackward) {
-    // y = [1, 0], ypred = [0.5, 0.5]
-    // diffs = [0.5, -0.5], MSE = 0.25, RMSE = 0.5
-    // grad_i = -(y_i - ypred_i) / (n * RMSE)
-    // grad[0] = -(1 - 0.5) / (2 * 0.5) = -0.5
-    // grad[1] = -(0 - 0.5) / (2 * 0.5) =  0.5
-    auto y = TensorFunctions::makeSharedTensor(
-        {2}, {1.0, 0.0}, false);
-    auto ypred = TensorFunctions::makeSharedTensor(
-        {2}, {0.5, 0.5}, true);
-
-    RmseLoss loss;
-    auto result = loss(y, ypred);
-    result->backward();
-
-    auto grads = ypred->getGrads();
-    EXPECT_NEAR((*grads)[0], -0.5f, kTol);
-    EXPECT_NEAR((*grads)[1],  0.5f, kTol);
+  // y = [1, 0], ypred = [0.5, 0.5]
+  // diffs = [0.5, -0.5], MSE = 0.25, RMSE = 0.5
+  // grad_i = -(y_i - ypred_i) / (n * RMSE)
+  // grad[0] = -(1 - 0.5) / (2 * 0.5) = -0.5
+  // grad[1] = -(0 - 0.5) / (2 * 0.5) =  0.5
+  auto y = TensorFunctions::makeSharedTensor(
+    {2}, {1.0, 0.0}, false);
+  auto ypred = TensorFunctions::makeSharedTensor(
+    {2}, {0.5, 0.5}, true);
+
+  RmseLoss loss;
+  auto result = loss(y, ypred);
+  result->backward();
+
+  auto grads = ypred->getGrads();
+  ASSERT_NEAR((*grads)[0], -0.5f, delta);
+  ASSERT_NEAR((*grads)[1],  0.5f, delta);
 }
\ No newline at end of file
diff --git a/tests/backend/test_module.cpp b/tests/backend/test_module.cpp
index a74c88c..ca0b1ac 100644
--- a/tests/backend/test_module.cpp
+++ b/tests/backend/test_module.cpp
@@ -23,21 +23,21 @@
 
 #include <cmath>
 
-constexpr ftype delta = 1e-3;
+using namespace std;
 
 TEST(ActivationTest, ReluForward) {
-  auto t1 = TensorFunctions::Ones({3, 2}, false);
+  auto t1 = TensorFunctions::Ones({3, 2});
   auto f = module::ReLu();
 
   auto res = f(t1);
 
   for(size_t i=0; i<t1.getSize(); i++){
-    ASSERT_DOUBLE_EQ(res[i], t1[i]);
+    ASSERT_NEAR(res[i], t1[i], 1e-5);
   }
 }
 
 TEST(ActivationTest, ReluInputNegative) {
-  auto t1 = TensorFunctions::Ones({3, 2}, false) * -1;
+  auto t1 = TensorFunctions::Ones({3, 2}) * -1;
   auto f = module::ReLu();
 
   auto res = f(t1);
@@ -49,191 +49,259 @@ TEST(ActivationTest, ReluInputNegative) {
 }
 
 TEST(AutogradTest, ReLUBackward) {
-    auto x = TensorFunctions::makeSharedTensor({3}, {-1.0, 0.0, 2.0}, true);
-    auto relu = module::ReLu();
+  auto x = TensorFunctions::makeSharedTensor({3}, {-1.0, 0.0, 2.0}, true);
+  auto relu = module::ReLu();
 
-    auto y = relu(x);    // [0, 0, 2]
-    auto loss = cgraph::sumTensor(y);  // loss = 2
-    
-    loss->backward();
-    
-    // Gradient: [0, 0, 1] (only where input > 0)
-    ASSERT_DOUBLE_EQ(x->getGrads()->get(0), 0.0);
-    ASSERT_DOUBLE_EQ(x->getGrads()->get(1), 0.0);
-    ASSERT_DOUBLE_EQ(x->getGrads()->get(2), 1.0);
+  auto y = relu(x);    // [0, 0, 2]
+  auto loss = cgraph::sumTensor(y);  // loss = 2
+
+  loss->backward();
+
+  // Gradient: [0, 0, 1] (only where input > 0)
+  ASSERT_DOUBLE_EQ(x->getGrads()->get(0), 0.0);
+  ASSERT_DOUBLE_EQ(x->getGrads()->get(1), 0.0);
+  ASSERT_NEAR(x->getGrads()->get(2), 1.0, 1e-5);
 }
 
 TEST(ActivationTest, LeakyReluForward) {
-  auto t1 = TensorFunctions::Ones({3, 2}, false);
+  auto t1 = TensorFunctions::Ones({3, 2});
 
   auto f = module::LeakyReLu(0.3);
   auto res = f(t1);
 
   for(size_t i=0; i<t1.getSize(); i++){
-    ASSERT_DOUBLE_EQ(res[i], t1[i]);
+    ASSERT_NEAR(res[i], t1[i], 1e-5);
   }
 }
 
 TEST(ActivationTest, LeakyReluInputNegative) {
-  auto t1 = TensorFunctions::Ones({3, 2}, false) * -1;
+  constexpr ftype factor = -1;
+  auto t1 = TensorFunctions::Ones({3, 2}) * factor;
   
   constexpr ftype eps = 0.3;
   auto f = module::LeakyReLu(eps);
   auto res = f(t1);
 
   for(size_t i=0; i<t1.getSize(); i++){
-    ASSERT_DOUBLE_EQ(res[i], eps);
+    ASSERT_NEAR(res[i], factor * eps, 1e-5);
   }
 }
 
 TEST(AutogradTest, LeakyReLUBackward) {
-    auto x = TensorFunctions::makeSharedTensor({3}, {-1.0, 0.0, 2.0}, true);
+  auto x = TensorFunctions::makeSharedTensor({3}, {-1.0, 0.0, 2.0}, true);
+
+  constexpr ftype eps = 0.3;
+  auto relu = module::LeakyReLu(eps);
 
-    constexpr ftype eps = 0.3;
-    auto relu = module::LeakyReLu(eps);
+  auto y = relu(x);    // [0, 0, 2]
+  auto loss = cgraph::sumTensor(y);  // loss = 2
 
-    auto y = relu(x);    // [0, 0, 2]
-    auto loss = cgraph::sumTensor(y);  // loss = 2
-    
-    loss->backward();
-    
-    // Gradient: [0, 0, 1] (only where input > 0)
-    ASSERT_DOUBLE_EQ(x->getGrads()->get(0), eps);
-    ASSERT_DOUBLE_EQ(x->getGrads()->get(1), eps); // by convention
-    ASSERT_DOUBLE_EQ(x->getGrads()->get(2), 1.0);
+  loss->backward();
+
+  // Gradient: [0, 0, 1] (only where input > 0)
+  ASSERT_DOUBLE_EQ(x->getGrads()->get(0), eps);
+  ASSERT_DOUBLE_EQ(x->getGrads()->get(1), eps); // by convention
+  ASSERT_NEAR(x->getGrads()->get(2), 1.0, 1e-5);
 }
 
 TEST(ActivationTest, SigmoidForward) {
-    // sigmoid(0) = 0.5, sigmoid(1) = 0.7311, sigmoid(-1) = 0.2689
-    auto t = Tensor({3}, {0.0, 1.0, -1.0}, true);
-    
-    module::Sigmoid sig;
-    auto res = sig(t);
+  // sigmoid(0) = 0.5, sigmoid(1) = 0.7311, sigmoid(-1) = 0.2689
+  auto t = Tensor({3}, {0.0, 1.0, -1.0});
 
-    EXPECT_NEAR(res[0], 0.5, delta);
-    EXPECT_NEAR(res[1], 0.7311, delta);
-    EXPECT_NEAR(res[2], 0.2689, delta);
+  module::Sigmoid sig;
+  auto res = sig(t);
+
+  ASSERT_NEAR(res[0], 0.5, 1e-4);
+  ASSERT_NEAR(res[1], 0.7311, 1e-4);
+  ASSERT_NEAR(res[2], 0.2689, 1e-4);
 }
 
 TEST(ActivationTest, SigmoidLargePositive) {
-    // sigmoid(100) should be ~1, not inf or nan
-    auto t = Tensor({1}, {100.0}, true);
-    
-    module::Sigmoid sig;
-    auto res = sig(t);
+  // sigmoid(100) should be ~1, not inf or nan
+  auto t = Tensor({1}, vector<ftype>{100.0});
+
+  module::Sigmoid sig;
+  auto res = sig(t);
 
-    EXPECT_NEAR(res[0], 1.0, delta);
-    EXPECT_FALSE(std::isnan(res[0]));
-    EXPECT_FALSE(std::isinf(res[0]));
+  ASSERT_NEAR(res[0], 1.0, 1e-5);
+  EXPECT_FALSE(std::isnan(res[0]));
+  EXPECT_FALSE(std::isinf(res[0]));
 }
 
 TEST(ActivationTest, SigmoidLargeNegative) {
-    // sigmoid(-100) should be ~0, not nan
-    auto t = Tensor({1}, {-100.0}, true);
-    
-    module::Sigmoid sig;
-    auto res = sig(t);
+  // sigmoid(-100) should be ~0, not nan
+  auto t = Tensor({1}, vector<ftype>{-100.0});
+
+  module::Sigmoid sig;
+  auto res = sig(t);
 
-    EXPECT_NEAR(res[0], 0.0, delta);
-    EXPECT_FALSE(std::isnan(res[0]));
-    EXPECT_FALSE(std::isinf(res[0]));
+  ASSERT_NEAR(res[0], 0.0, 1e-5);
+  EXPECT_FALSE(std::isnan(res[0]));
+  EXPECT_FALSE(std::isinf(res[0]));
 }
 
 TEST(AutogradTest, SigmoidBackward) {
-    // grad of sigmoid = sigmoid(x) * (1 - sigmoid(x))
-    // for x=0: grad = 0.5 * 0.5 = 0.25
-    // for x=1: grad = 0.7311 * 0.2689 = 0.1966
-    auto t = TensorFunctions::makeSharedTensor(
-        {2}, {0.0, 1.0}, true);
-    
-    module::Sigmoid sig;
-    auto res = sig(t);
-    res->backward();
-
-    auto grads = t->getGrads();
-    EXPECT_NEAR((*grads)[0], 0.25, delta);
-    EXPECT_NEAR((*grads)[1], 0.1966, delta);
+  // grad of sigmoid = sigmoid(x) * (1 - sigmoid(x))
+  // for x=0: grad = 0.5 * 0.5 = 0.25
+  // for x=1: grad = 0.7311 * 0.2689 = 0.1966
+  auto t = TensorFunctions::makeSharedTensor(
+    {2}, {0.0, 1.0}, true);
+
+  module::Sigmoid sig;
+  auto res = sig(t);
+  res->backward();
+
+  auto grads = t->getGrads();
+  ASSERT_NEAR((*grads)[0], 0.25, 1e-4);
+  ASSERT_NEAR((*grads)[1], 0.1966, 1e-4);
 }
 
 TEST(ActivationTest, SoftmaxForward) {
-    // softmax([1, 2, 3])
-    // exp([1,2,3]) = [2.7183, 7.3891, 20.0855]
-    // sum = 30.1929
-    // softmax = [0.0900, 0.2447, 0.6652]
-    auto t = Tensor({1, 3}, {1.0, 2.0, 3.0}, true);
-    
-    module::Softmax sm;
-    auto res = sm(t);
-
-    EXPECT_NEAR(res[0], 0.0900, delta);
-    EXPECT_NEAR(res[1], 0.2447, delta);
-    EXPECT_NEAR(res[2], 0.6652, delta);
+  // softmax([1, 2, 3])
+  // exp([1,2,3]) = [2.7183, 7.3891, 20.0855]
+  // sum = 30.1929
+  // softmax = [0.0900, 0.2447, 0.6652]
+  auto t = Tensor({1, 3}, {1.0, 2.0, 3.0});
+
+  module::Softmax sm;
+  auto res = sm(t);
+
+  ASSERT_NEAR(res[0], 0.0900, 1e-4);
+  ASSERT_NEAR(res[1], 0.2447, 1e-4);
+  ASSERT_NEAR(res[2], 0.6652, 1e-4);
 }
 
 TEST(ActivationTest, SoftmaxSumsToOne) {
-    auto t = Tensor({2, 4}, 
-                    {1.0, 2.0, 3.0, 4.0,
-                     2.0, 1.0, 4.0, 3.0}, 
-                     true);
-    
-    module::Softmax sm;
-    auto res = sm(t);
-
-    // each row must sum to 1
-    ftype row0sum = res[0] + res[1] + res[2] + res[3];
-    ftype row1sum = res[4] + res[5] + res[6] + res[7];
-    EXPECT_NEAR(row0sum, 1.0, delta);
-    EXPECT_NEAR(row1sum, 1.0, delta);
+  auto t = Tensor({2, 4},
+                  {1.0, 2.0, 3.0, 4.0,
+                   2.0, 1.0, 4.0, 3.0});
+
+  module::Softmax sm;
+  auto res = sm(t);
+
+  // each row must sum to 1
+  ftype row0sum = res[0] + res[1] + res[2] + res[3];
+  ftype row1sum = res[4] + res[5] + res[6] + res[7];
+  ASSERT_NEAR(row0sum, 1.0, 1e-5);
+  ASSERT_NEAR(row1sum, 1.0, 1e-5);
 }
 
 TEST(ActivationTest, SoftmaxForwardNumericalStability) {
-    // large values should not produce nan or inf
-    auto t = Tensor({1, 3}, {100.0, 101.0, 102.0}, true);
-    
-    module::Softmax sm;
-    auto res = sm(t);
+  // large values should not produce nan or inf
+  auto t = Tensor({1, 3}, {100.0, 101.0, 102.0});
 
-    for(int i = 0; i < 3; i++) {
-        EXPECT_FALSE(std::isnan(res[i]));
-        EXPECT_FALSE(std::isinf(res[i]));
-    }
-    ftype rowsum = res[0] + res[1] + res[2];
-    EXPECT_NEAR(rowsum, 1.0, delta);
+  module::Softmax sm;
+  auto res = sm(t);
+
+  for(int i = 0; i < 3; i++) {
+    EXPECT_FALSE(std::isnan(res[i]));
+    EXPECT_FALSE(std::isinf(res[i]));
+  }
+
+  ftype rowsum = res[0] + res[1] + res[2];
+  ASSERT_NEAR(rowsum, 1.0, 1e-5);
 }
 
 TEST(AutogradTest, SoftmaxBackward) {
-    // for softmax with upstream grad of ones, the gradient is zero
-    // because d/dx_i sum(softmax(x)) = 0 (softmax sums to 1 always)
-    // more useful: upstream = [1, 0, 0]
-    // grad[i] = softmax[i] * (upstream[i] - dot(upstream, softmax))
-    // for x=[1,2,3], softmax=[0.09, 0.2447, 0.6652]
-    // dot([1,0,0], softmax) = 0.09
-    // grad[0] = 0.09   * (1 - 0.09)   =  0.0819
-    // grad[1] = 0.2447 * (0 - 0.09)   = -0.0220
-    // grad[2] = 0.6652 * (0 - 0.09)   = -0.0599
-    auto t = TensorFunctions::makeSharedTensor(
-        {1, 3}, {1.0, 2.0, 3.0}, true);
+  // for softmax with upstream grad of ones, the gradient is zero
+  // because d/dx_i sum(softmax(x)) = 0 (softmax sums to 1 always)
+  // more useful: upstream = [1, 0, 0]
+  // grad[i] = softmax[i] * (upstream[i] - dot(upstream, softmax))
+  // for x=[1,2,3], softmax=[0.09, 0.2447, 0.6652]
+  // dot([1,0,0], softmax) = 0.09
+  // grad[0] = 0.09   * (1 - 0.09)   =  0.0819
+  // grad[1] = 0.2447 * (0 - 0.09)   = -0.0220
+  // grad[2] = 0.6652 * (0 - 0.09)   = -0.0599
+  auto t = TensorFunctions::makeSharedTensor({1, 3}, {1.0, 2.0, 3.0}, true);
     
-    module::Softmax sm;
-    auto resPtr = sm(t);
+  module::Softmax sm;
+  auto resPtr = sm(t);
     
-    // set upstream gradient to [1, 0, 0]
-    auto upstream = TensorFunctions::makeSharedTensor(
-        {1, 3}, {1.0, 0.0, 0.0}, false);
-    resPtr->setGrads(upstream);
-    resPtr->backward();
-
-    auto grads = t->getGrads();
-    EXPECT_NEAR((*grads)[0],  0.0819, delta);
-    EXPECT_NEAR((*grads)[1], -0.0220, delta);
-    EXPECT_NEAR((*grads)[2], -0.0599, delta);
+  // set upstream gradient to [1, 0, 0]
+  auto upstream = TensorFunctions::makeSharedTensor(
+    {1, 3}, {1.0, 0.0, 0.0});
+  resPtr->setGrads(upstream);
+  resPtr->backward();
+
+  auto grads = t->getGrads();
+  ASSERT_NEAR((*grads)[0],  0.0819, 1e-4);
+  ASSERT_NEAR((*grads)[1], -0.0220, 1e-4);
+  ASSERT_NEAR((*grads)[2], -0.0599, 1e-4);
 }
 
 TEST(LayerTest, TestFfLayer) {
-  auto t1 = TensorFunctions::Ones({3, 2}, false);
-  auto layer = module::FfLayer(2, 1, true, false);
+  auto t1 = TensorFunctions::Ones({3, 2});
+  auto layer = module::FfLayer(2, 1);
 
   auto res = layer(t1);
 
   ASSERT_EQ(res.getDims(), Dimension({3, 1}));
+}
+
+TEST(AutogradTest, FfLayerBackward) {
+  // x: (2,3) all ones, W: (3,2) all ones
+  // y = x @ W => (2,2) all threes
+  // loss = sum(y), upstream = ones(2,2)
+  // grad_x = upstream @ W^T = ones(2,2) @ ones(2,3) => all 2s
+  // grad_W = x^T @ upstream = ones(3,2) @ ones(2,2) => all 2s
+  auto x = TensorFunctions::makeSharedTensor({2, 3}, {
+    1.0, 1.0, 1.0,
+    1.0, 1.0, 1.0
+  }, true);
+
+  auto layer = module::FfLayer(3, 2, /*useBias=*/false, /*requiresGrad=*/true);
+  auto w = layer.getWeights();
+  for(tensorSize_t i = 0; i < w->getSize(); i++) w->set(1.0, i);
+
+  auto res = layer(x);
+  auto loss = cgraph::sumTensor(res);
+  loss->backward();
+
+  auto xGrads = x->getGrads();
+  ASSERT_NE(xGrads, nullptr);
+  ASSERT_EQ(xGrads->getDims(), x->getDims());
+  for(tensorSize_t i = 0; i < xGrads->getSize(); i++) {
+    ASSERT_NEAR((*xGrads)[i], 2.0, 1e-5);
+  }
+
+  auto wGrads = layer.getWeights()->getGrads();
+  ASSERT_NE(wGrads, nullptr);
+  ASSERT_EQ(wGrads->getDims(), layer.getWeights()->getDims());
+  for(tensorSize_t i = 0; i < wGrads->getSize(); i++) {
+    ASSERT_NEAR((*wGrads)[i], 2.0, 1e-5);
+  }
+}
+
+TEST(AutogradTest, FfLayerBackwardWithBias) {
+  // Same as above but with bias (initialised to zero by constructor).
+  // grad_b = sum(upstream, axis=0) = [2, 2]
+  auto x = TensorFunctions::makeSharedTensor({2, 3}, {
+    1.0, 1.0, 1.0,
+    1.0, 1.0, 1.0
+  }, true);
+
+  auto layer = module::FfLayer(3, 2, /*useBias=*/true, /*requiresGrad=*/true);
+  auto w = layer.getWeights();
+  for(tensorSize_t i = 0; i < w->getSize(); i++) w->set(1.0, i);
+
+  auto res = layer(x);
+  auto loss = cgraph::sumTensor(res);
+  loss->backward();
+
+  auto xGrads = x->getGrads();
+  ASSERT_NE(xGrads, nullptr);
+  for(tensorSize_t i = 0; i < xGrads->getSize(); i++) {
+    ASSERT_NEAR((*xGrads)[i], 2.0, 1e-5);
+  }
+
+  auto wGrads = layer.getWeights()->getGrads();
+  ASSERT_NE(wGrads, nullptr);
+  for(tensorSize_t i = 0; i < wGrads->getSize(); i++) {
+    ASSERT_NEAR((*wGrads)[i], 2.0, 1e-5);
+  }
+
+  auto bGrads = layer.getBias()->getGrads();
+  ASSERT_NE(bGrads, nullptr);
+  ASSERT_NEAR((*bGrads)[0], 2.0, 1e-5);
+  ASSERT_NEAR((*bGrads)[1], 2.0, 1e-5);
 }
\ No newline at end of file
diff --git a/tests/backend/test_tensorops.cpp b/tests/backend/test_tensorops.cpp
new file mode 100644
index 0000000..0a197e3
--- /dev/null
+++ b/tests/backend/test_tensorops.cpp
@@ -0,0 +1,354 @@
+/**
+ * @file test_data_modeling.cpp
+ * @author Robert Baumgartner (r.baumgartner-1@tudelft.nl)
+ * @brief 
+ * @version 0.1
+ * @date 2026-02-16
+ * 
+ * @copyright Copyright (c) 2026
+ * 
+ */
+
+#include <gtest/gtest.h>
+
+#include "data_modeling/tensor.h"
+#include "data_modeling/tensor_functions.h"
+
+#include "computational_graph/tensor_ops/graph_creation.h"
+
+using namespace std;
+
+TEST(TensorOpsTest, TestCtor) {
+  auto t = Tensor({2, 2}, {2.0, 3.0, 4.0, 5.0}, Device::CPU);
+
+  ASSERT_EQ(t.getDims(), Dimension({2, 2}));
+  ASSERT_EQ(t.getDevice(), Device::CPU);
+  ASSERT_TRUE(!t.getRequiresGrad());
+
+  ASSERT_NEAR(t.get(0, 0), 2.0, 1e-5);
+  ASSERT_NEAR(t.get(0, 1), 3.0, 1e-5);
+  ASSERT_NEAR(t.get(1, 0), 4.0, 1e-5);
+  ASSERT_NEAR(t.get(1, 1), 5.0, 1e-5);
+}
+
+TEST(TensorOpsTest, ScalarAdd) {
+  auto t1 = TensorFunctions::Ones({2, 2});
+
+  auto res = t1 + 1.5;
+
+  constexpr ftype sum = 2.5;
+  for(auto i=0; i<t1.getDims().get(0); i++) {
+    for(auto j=0; j<t1.getDims().get(1); j++) {
+      ASSERT_NEAR(res.get(i, j), sum, 1e-5);
+    }
+  }
+}
+
+TEST(TensorOpsTest, ScalarMul) {
+  auto t1 = TensorFunctions::Ones({2, 2});
+
+  constexpr ftype f = 2.5;
+  auto res = t1 * f;
+    
+  for(auto i=0; i<t1.getDims().get(0); i++) {
+    for(auto j=0; j<t1.getDims().get(1); j++) {
+      ASSERT_NEAR(res.get(i, j), f, 1e-5);
+    }
+  }
+}
+
+TEST(AutogradTest, ScalarMul) {
+  auto t1 = TensorFunctions::makeSharedTensor({1}, {2.0}, true);
+  auto t2 = TensorFunctions::makeSharedTensor({1}, {3.0}, true);
+
+  auto t3 = cgraph::mul(t1, t2);
+  auto loss = cgraph::mul(t3, t3);
+
+  loss->backward();
+
+  ASSERT_DOUBLE_EQ(t1->getGrads()->get(0), 36.0);
+  ASSERT_DOUBLE_EQ(t2->getGrads()->get(0), 24.0);
+}
+
+TEST(TensorOpsTest, TensorAdd) {
+  auto t1 = TensorFunctions::Ones({2, 2});
+  auto t2 = TensorFunctions::Ones({2, 2}) * 4;
+
+  auto res = t1 + t2;
+
+  constexpr ftype sum = 5.0;
+  for(auto i=0; i<t1.getDims().get(0); i++) {
+    for(auto j=0; j<t1.getDims().get(1); j++) {
+      ASSERT_NEAR(res.get(i, j), sum, 1e-5);
+    }
+  }
+}
+
+TEST(TensorOpsTest, TensorAddThrowsOnDimMismatch) {
+  auto t1 = TensorFunctions::Ones({2, 2});
+  auto t2 = TensorFunctions::Ones({2, 3}) * 4;
+
+  ASSERT_THROW(t1 + t2, std::invalid_argument);
+}
+
+TEST(AutogradTest, TensorAdd) {
+  auto t1 = TensorFunctions::makeSharedTensor({1}, {3.0}, true);
+  auto t2 = TensorFunctions::makeSharedTensor({1}, {2.0}, true);
+
+  auto t3 = cgraph::add(t1, t2);
+  auto loss = cgraph::mul(t3, t3);
+
+  loss->backward();
+
+  ASSERT_NEAR(t1->getGrads()->get(0), 10.0, 1e-5);
+  ASSERT_NEAR(t2->getGrads()->get(0), 10.0, 1e-5);
+}
+
+TEST(TensorOpsTest, BroadcastAdd) {
+  auto t1 = TensorFunctions::Ones({3, 2, 2});
+  auto t2 = Tensor({2}, {2, 3});
+
+  auto res = t1 + t2;
+
+  ASSERT_EQ(res.getDims(), t1.getDims());
+  
+  for(auto i=0; i<res.getDims().get(0); i++) {
+    for(auto j=0; j<res.getDims().get(1); j++) {
+      ASSERT_NEAR(res.get(i, j, 0), 3.0, 1e-5);
+      ASSERT_NEAR(res.get(i, j, 1), 4.0, 1e-5);
+    }
+  }
+}
+
+TEST(TensorOpsTest, BroadcastAdd2D) {
+  // (2,3) + (3)
+  auto t1 = Tensor({2, 3}, {1.0, 2.0, 3.0,
+                            4.0, 5.0, 6.0});
+  auto t2 = Tensor({3}, {10.0, 20.0, 30.0});
+  auto res = t1 + t2;
+
+  // expected: each row of t1 gets t2 added elementwise
+  ASSERT_NEAR(res.get(0, 0), 11.0, 1e-5);
+  ASSERT_NEAR(res.get(0, 1), 22.0, 1e-5);
+  ASSERT_NEAR(res.get(0, 2), 33.0, 1e-5);
+  ASSERT_NEAR(res.get(1, 0), 14.0, 1e-5);
+  ASSERT_NEAR(res.get(1, 1), 25.0, 1e-5);
+  ASSERT_NEAR(res.get(1, 2), 36.0, 1e-5);
+}
+
+TEST(TensorOpsTest, BroadcastAddNotCommutative) {
+  auto t1 = TensorFunctions::Ones({3, 2, 2});
+  auto t2 = Tensor({2}, {2, 3});
+
+  ASSERT_THROW(t2 + t1, std::invalid_argument);
+}
+
+TEST(AutogradTest, BroadcastAdd) {
+  // gradient of broadcast add w.r.t. bias should be sum over batch dimension
+  // upstream grad: (2,3) of ones → bias grad should be (3) of twos
+  auto t1 = TensorFunctions::makeSharedTensor({2, 3},
+    {1.0, 2.0, 3.0,
+     4.0, 5.0, 6.0}, true);
+  auto bias = TensorFunctions::makeSharedTensor({3},
+    {0.0, 0.0, 0.0}, true);
+
+  auto res = cgraph::add(t1, bias);
+  res->backward();
+
+  // bias grad should be sum over batch: [2, 2, 2]
+  auto biasGrad = bias->getGrads();
+  ASSERT_NEAR((*biasGrad)[0], 2.0, 1e-5);
+  ASSERT_NEAR((*biasGrad)[1], 2.0, 1e-5);
+  ASSERT_NEAR((*biasGrad)[2], 2.0, 1e-5);
+
+  // t1 grad should be ones (add is identity for non-broadcast operand)
+  auto t1Grad = t1->getGrads();
+  for(int i = 0; i < 6; i++) {
+    ASSERT_NEAR((*t1Grad)[i], 1.0, 1e-5);
+  }
+}
+
+TEST(TensorOpsTest, MatrixAdd) {
+  auto t1 = TensorFunctions::Ones({2, 2});
+  auto t2 = TensorFunctions::Ones({2, 2});
+    
+  auto res = t1 + t2;
+    
+  constexpr ftype resSum = 2.0;
+
+  for(auto i=0; i<t1.getDims().get(0); i++) {
+    for(auto j=0; j<t1.getDims().get(1); j++) {
+      ASSERT_NEAR(res.get(i, j), resSum, 1e-5);
+    }
+  }
+}
+
+TEST(TensorOpsTest, ElementwiseMul) {
+  constexpr ftype factor = 0.5;
+  auto t1 = TensorFunctions::Ones({2, 2});
+  auto t2 = TensorFunctions::Ones({2, 2}) * 0.5;
+    
+  auto res = t1 * t2;
+    
+  for(auto i=0; i<t1.getDims().get(0); i++) {
+    for(auto j=0; j<t1.getDims().get(1); j++) {
+      ASSERT_NEAR(res.get(i, j), factor, 1e-5);
+    }
+  }
+}
+
+TEST(TensorOpsTest, ElementwiseMulThrowsOnDimensionMismatch) {
+  constexpr ftype factor = 0.5;
+  auto t1 = TensorFunctions::Ones({2, 2});
+  auto t2 = TensorFunctions::Ones({2, 3}) * 0.5;
+    
+  ASSERT_THROW(t1 * t2, std::invalid_argument);
+}
+
+TEST(TensorOpsTest, MatMul) {
+  auto t1 = TensorFunctions::Ones({3, 2});
+  auto t2 = TensorFunctions::Ones({2, 6}) * 1.5;
+    
+  auto res = t1.matmul(t2);
+  auto expectedDims = std::vector<tensorDim_t>{3, 6};
+  ASSERT_EQ(res.getDims().toVector(), expectedDims);
+
+  constexpr ftype resSum = 3.0;
+  for(auto i=0; i<t1.getDims().get(0); i++) {
+    for(auto j=0; j<t1.getDims().get(1); j++) {
+      ASSERT_NEAR(res.get(i, j), resSum, 1e-5);
+    }
+  }
+}
+
+TEST(TensorOpsTest, MatMul2) {
+  auto t1 = Tensor({2, 2});
+  auto t2 = Tensor({2, 2});
+
+  auto cmpRes = Tensor({2, 2});
+
+  auto populateTensor = [](Tensor& t, ftype v1, ftype v2, ftype v3, ftype v4) {
+    t.set(v1, {0, 0});
+    t.set(v2, {0, 1});
+    t.set(v3, {1, 0});
+    t.set(v4, {1, 1});
+  };
+
+  populateTensor(t1, 1, 2, 3, 4);
+  populateTensor(t2, 5, 6, 7, 8);
+  populateTensor(cmpRes, 19, 22, 43, 50);
+
+  auto res = t1.matmul(t2);
+
+  auto expectedDims = std::vector<tensorDim_t>{2, 2};
+  ASSERT_EQ(res.getDims().toVector(), expectedDims);
+
+  constexpr ftype resSum = 3.0;
+  for(auto i=0; i<t1.getDims().get(0); i++) {
+    for(auto j=0; j<t1.getDims().get(1); j++) {
+      ASSERT_NEAR(res.get(i, j), cmpRes.get(i, j), 1e-5);
+    }
+  }
+}
+
+TEST(TensorOpsTest, MatMulThrowsWhenDimensionsNotMatched) {
+  auto t1 = TensorFunctions::Ones({2, 2});
+  auto t2 = TensorFunctions::Ones({3, 2});
+
+  ASSERT_THROW(t1.matmul(t2), std::runtime_error);
+}
+
+TEST(AutogradTest, MatMul) {
+  auto t1 = TensorFunctions::makeSharedTensor({2, 3}, {1, 2, 3, 4, 5, 6}, true);
+  auto t2 = TensorFunctions::makeSharedTensor({3, 2}, {1, 2, 3, 4, 5, 6}, true);
+
+  auto t3 = cgraph::matmul(t1, t2);
+
+  auto loss = TensorFunctions::makeSharedTensor({1}, {0.0}, true);
+  for(size_t i = 0; i < t3->getSize(); ++i) {
+    loss = cgraph::add(loss, cgraph::get(t3, i));
+  }
+
+  loss->backward();
+
+  EXPECT_TRUE(t1->hasGrads());
+  EXPECT_TRUE(t2->hasGrads());
+
+  // dL/dt1 = dloss/dt3 @ t2^t = Ones({2, 2}) @ t2^t
+  ASSERT_DOUBLE_EQ(t1->getGrads()->get({0, 0}), 3.0);
+  ASSERT_DOUBLE_EQ(t1->getGrads()->get({0, 1}), 7.0);
+  ASSERT_DOUBLE_EQ(t1->getGrads()->get({0, 2}), 11.0);
+  ASSERT_DOUBLE_EQ(t1->getGrads()->get({1, 0}), 3.0);
+  ASSERT_DOUBLE_EQ(t1->getGrads()->get({1, 1}), 7.0);
+  ASSERT_DOUBLE_EQ(t1->getGrads()->get({1, 2}), 11.0);
+
+  // dL/dt2 = t1^t @ dloss/dt3 = t1^t @ Ones({2, 2})
+  ASSERT_DOUBLE_EQ(t2->getGrads()->get({0, 0}), 5.0);
+  ASSERT_DOUBLE_EQ(t2->getGrads()->get({0, 1}), 5.0);
+  ASSERT_DOUBLE_EQ(t2->getGrads()->get({1, 0}), 7.0);
+  ASSERT_DOUBLE_EQ(t2->getGrads()->get({1, 1}), 7.0);
+  ASSERT_DOUBLE_EQ(t2->getGrads()->get({2, 0}), 9.0);
+  ASSERT_DOUBLE_EQ(t2->getGrads()->get({2, 1}), 9.0);
+}
+
+
+TEST(TensorOpsTest, MatrixTranspose) {
+  auto t = TensorFunctions::Gaussian({3, 2}, 1.0);
+
+  auto transposed = t.createDeepCopy();
+  transposed = transposed.transpose(-1, -2);
+
+  ASSERT_EQ(t.getDims().get(-1), transposed.getDims().get(-2));
+  ASSERT_EQ(t.getDims().get(-2), transposed.getDims().get(-1));
+  ASSERT_EQ(t.getDims().nDims(), transposed.getDims().nDims());
+  
+  for(auto row=0; row<t.getDims().get(-2); row++) {
+    for(auto col=0; col<t.getDims().get(-1); col++) {
+      ASSERT_NEAR(t.get(row, col), transposed.get(col, row), 1e-5);
+    }
+  }
+}
+
+/**
+ * @brief Swap first two dimensions.
+ */
+TEST(TensorOpsTest, MatrixTranspose2) {
+  auto t = TensorFunctions::Gaussian({3, 2, 5}, 1.0);
+  auto transposed = t.createDeepCopy().transpose(0, 1);
+
+  ASSERT_EQ(t.getDims().get(0), transposed.getDims().get(1));
+  ASSERT_EQ(t.getDims().get(1), transposed.getDims().get(0));
+  ASSERT_EQ(t.getDims().get(-1), transposed.getDims().get(-1));
+  ASSERT_EQ(t.getDims().nDims(), transposed.getDims().nDims());
+  
+  for(auto dim1=0; dim1<t.getDims().get(0); dim1++) {
+    for(auto dim2=0; dim2<t.getDims().get(1); dim2++) {
+      for(auto dim3=0; dim3<t.getDims().get(-1); dim3++) {
+        // we transposed dim1 and dim3
+        ASSERT_NEAR(t.get(dim1, dim2, dim3), transposed.get(dim2, dim1, dim3), 1e-5);
+      }
+    }
+  }
+}
+
+/**
+ * @brief Swap first and last dimension.
+ */
+TEST(TensorOpsTest, MatrixTranspose3) {
+  auto t = TensorFunctions::Gaussian({3, 2, 5}, 1.0);
+  auto transposed = t.createDeepCopy().transpose(0, -1);
+
+  ASSERT_EQ(t.getDims().get(0), transposed.getDims().get(-1));
+  ASSERT_EQ(t.getDims().get(-1), transposed.getDims().get(0));
+  ASSERT_EQ(t.getDims().get(1), transposed.getDims().get(1));
+  ASSERT_EQ(t.getDims().nDims(), transposed.getDims().nDims());
+  
+  for(auto dim1=0; dim1<t.getDims().get(0); dim1++) {
+    for(auto dim2=0; dim2<t.getDims().get(1); dim2++) {
+      for(auto dim3=0; dim3<t.getDims().get(-1); dim3++) {
+        // we transposed dim1 and dim3
+        ASSERT_NEAR(t.get(dim1, dim2, dim3), transposed.get(dim3, dim2, dim1), 1e-5);
+      }
+    }
+  }
+}
diff --git a/tests/backend/test_train_loop.cpp b/tests/backend/test_train_loop.cpp
deleted file mode 100644
index 4ed7b9d..0000000
--- a/tests/backend/test_train_loop.cpp
+++ /dev/null
@@ -1,278 +0,0 @@
-/**
- * @file test_training.cpp
- * @author Robert Baumgartner (r.baumgartner-1@tudelft.nl)
- * @brief 
- * @version 0.1
- * @date 2026-03-14
- * 
- * @copyright Copyright (c) 2026
- * 
- */
-
-#include <gtest/gtest.h>
-
-#include "module/networks/sequential.h"
-#include "module/layers/ff_layer.h"
-
-#include "module/activation_functions/sigmoid.h"
-#include "module/activation_functions/relu.h"
-#include "module/activation_functions/leaky_relu.h"
-#include "module/activation_functions/softmax.h"
-
-#include "training/optimizers/sgd.h"
-#include "training/optimizers/rmsprop.h"
-
-#include "training/loss_functions/bce_loss.h"
-#include "training/loss_functions/crossentropy_loss.h"
-#include "training/loss_functions/bce_sigmoid_loss.h"
-#include "training/loss_functions/crossentropy_softmax_loss.h"
-
-#include "training/trainers/base_train_loop.h"
-
-#include "data_modeling/tensor_functions.h"
-
-#include "system/sys_functions.h"
-
-using namespace std;
-
-static shared_ptr<module::Sequential> makeBinaryNet() {
-    auto net = make_shared<module::Sequential>();
-
-    net->append(make_shared<module::FfLayer>(2, 4, true, true));
-
-    net->append(make_shared<module::LeakyReLu>(0.01));
-
-    net->append(make_shared<module::FfLayer>(4, 1, true, true));
-
-    net->append(make_shared<module::Sigmoid>());
-    return net;
-}
-
-static shared_ptr<module::Sequential> makeBinaryNet2() {
-    auto net = make_shared<module::Sequential>();
-
-    net->append(make_shared<module::FfLayer>(2, 4, true, true));
-
-    net->append(make_shared<module::LeakyReLu>(0.01));
-
-    net->append(make_shared<module::FfLayer>(4, 1, true, true));
-
-    return net;
-}
-
-static shared_ptr<module::Sequential> makeMulticlassNet() {
-    auto net = make_shared<module::Sequential>();
-
-    net->append(make_shared<module::FfLayer>(2, 8, true, true));
-
-    net->append(make_shared<module::LeakyReLu>(0.01));
-
-    net->append(make_shared<module::FfLayer>(8, 3, true, true));
-
-    net->append(make_shared<module::Softmax>());
-    return net;
-}
-
-static shared_ptr<module::Sequential> makeMulticlassNet2() {
-    auto net = make_shared<module::Sequential>();
-
-    net->append(make_shared<module::FfLayer>(2, 8, true, true));
-
-    net->append(make_shared<module::LeakyReLu>(0.01));
-
-    net->append(make_shared<module::FfLayer>(8, 3, true, true));
-
-    return net;
-}
-
-int main(int argc, char** argv) {
-    testing::InitGoogleTest(&argc, argv);
-    sys::setRandomSeed(42);
-    return RUN_ALL_TESTS();
-}
-
-TEST(OverfitTest, BceSgdOverfitsSmallDataset) {
-    // XOR-like: 4 samples, 2 features, binary labels
-    auto x = TensorFunctions::makeSharedTensor(
-        {4, 2}, {0.0, 0.0,
-                 0.0, 1.0,
-                 1.0, 0.0,
-                 1.0, 1.0}, false);
-
-    auto y = TensorFunctions::makeSharedTensor(
-        {4, 1}, {0.0,
-                 1.0,
-                 1.0,
-                 0.0}, false);
-
-    auto net = makeBinaryNet();    
-    auto loss = make_shared<train::BceLoss>();
-    auto optim = make_shared<train::SgdOptimizer>(
-        net->parameters(), /*lr=*/0.05);
-
-    auto trainLoop = train::BaseTrainLoop(
-        net, loss, optim, /*epochs=*/2000, /*bsize=*/static_cast<tensorDim_t>(4));
-
-    trainLoop.run(x, y, /*shuffle=*/false, /*verbose=*/false);
-
-    // forward one more time to get final loss
-    auto pred = (*net)(x);
-    auto finalLoss = (*loss)(y, pred);
-
-    EXPECT_LT((*finalLoss)[0], 0.05f)
-        << "Network failed to overfit binary dataset\n"
-        << "Final prediction: " << *pred << "\nFinal loss: " << *finalLoss;
-}
-
-TEST(OverfitTest, BceSgdOverfitsSmallDataset_OptimizedLoss) {    
-    // XOR-like: 4 samples, 2 features, binary labels
-    auto x = TensorFunctions::makeSharedTensor(
-        {4, 2}, {0.0, 0.0,
-                 0.0, 1.0,
-                 1.0, 0.0,
-                 1.0, 1.0}, false);
-
-    auto y = TensorFunctions::makeSharedTensor(
-        {4, 1}, {0.0,
-                 1.0,
-                 1.0,
-                 0.0}, false);
-
-    auto net = makeBinaryNet2();    
-    auto loss = make_shared<train::BceSigmoidLoss>();
-    auto optim = make_shared<train::SgdOptimizer>(
-        net->parameters(), /*lr=*/0.05);
-
-    auto trainLoop = train::BaseTrainLoop(
-        net, loss, optim, /*epochs=*/2000, /*bsize=*/static_cast<tensorDim_t>(4));
-
-    trainLoop.run(x, y, /*shuffle=*/false, /*verbose=*/false);
-
-    // forward one more time to get final loss
-    auto pred = (*net)(x);
-    auto finalLoss = (*loss)(y, pred);
-
-    auto sigmoid = module::Sigmoid();
-    EXPECT_LT((*finalLoss)[0], 0.05f)
-        << "Network failed to overfit binary dataset\n"
-        << "Final prediction: " << sigmoid(*pred) << "\nFinal loss: " << *finalLoss;
-}
-
-TEST(OverfitTest, CrossEntropyRMSPropOverfitsSmallDataset) {
-    // 6 samples, 2 features, 3 classes
-    auto x = TensorFunctions::makeSharedTensor(
-        {6, 2}, {1.0, 0.0,
-                 1.0, 0.1,
-                 0.0, 1.0,
-                 0.1, 1.0,
-                 0.5, 0.5,
-                 0.4, 0.6}, false);
-
-    // one-hot encoded labels
-    auto y = TensorFunctions::makeSharedTensor(
-        {6, 3}, {1.0, 0.0, 0.0,
-                 1.0, 0.0, 0.0,
-                 0.0, 1.0, 0.0,
-                 0.0, 1.0, 0.0,
-                 0.0, 0.0, 1.0,
-                 0.0, 0.0, 1.0}, false);
-
-    auto net = makeMulticlassNet();
-    auto loss = make_shared<train::CrossEntropyLoss>();
-    auto optim = make_shared<train::RmsPropOptimizer>(
-        net->parameters(), /*lr=*/0.0001, /*decay=*/0.95);
-
-    auto trainLoop = train::BaseTrainLoop(
-        net, loss, optim, /*epochs=*/2000, /*bsize=*/6);
-
-    trainLoop.run(x, y, /*shuffle=*/false, /*verbose=*/false);
-
-    auto pred = (*net)(x);
-    auto finalLoss = (*loss)(y, pred);
-
-    EXPECT_LT((*finalLoss)[0], 0.05f)
-        << "Network failed to overfit multiclass dataset"
-        << "Final prediction: " << *pred << "\nFinal loss: " << *finalLoss;
-}
-
-TEST(OverfitTest, CrossEntropyRMSPropOverfitsSmallDataset_OptimizedLoss) {
-    // 6 samples, 2 features, 3 classes
-    auto x = TensorFunctions::makeSharedTensor(
-        {6, 2}, {1.0, 0.0,
-                 1.0, 0.1,
-                 0.0, 1.0,
-                 0.1, 1.0,
-                 0.5, 0.5,
-                 0.4, 0.6}, false);
-
-    // one-hot encoded labels
-    auto y = TensorFunctions::makeSharedTensor(
-        {6, 3}, {1.0, 0.0, 0.0,
-                 1.0, 0.0, 0.0,
-                 0.0, 1.0, 0.0,
-                 0.0, 1.0, 0.0,
-                 0.0, 0.0, 1.0,
-                 0.0, 0.0, 1.0}, false);
-
-    auto net = makeMulticlassNet2();
-    auto loss = make_shared<train::CrossEntropySoftmaxLoss>();
-    auto optim = make_shared<train::RmsPropOptimizer>(
-        net->parameters(), /*lr=*/0.0003, /*decay=*/0.95);
-
-    auto trainLoop = train::BaseTrainLoop(
-        net, loss, optim, /*epochs=*/10000, /*bsize=*/6);
-
-    trainLoop.run(x, y, /*shuffle=*/false, /*verbose=*/false);
-
-    auto pred = (*net)(x);
-    auto finalLoss = (*loss)(y, pred);
-
-    auto softmax = module::Softmax();
-    EXPECT_LT((*finalLoss)[0], 0.05f)
-        << "Network failed to overfit multiclass dataset"
-        << "Final prediction: " << softmax(*pred) << "\nFinal loss: " << *finalLoss;
-}
-
-TEST(OptimizerTest, ZeroGrad_ClearsAllGradients) {
-    auto x = TensorFunctions::makeSharedTensor(
-        {4, 2}, {0.0, 0.0, 0.0, 1.0, 1.0, 0.0, 1.0, 1.0}, false);
-    auto y = TensorFunctions::makeSharedTensor(
-        {4, 1}, {0.0, 1.0, 1.0, 0.0}, false);
-
-    auto net = makeBinaryNet();
-    auto loss = std::make_shared<train::BceSigmoidLoss>();
-    auto optim = std::make_shared<train::SgdOptimizer>(
-        net->parameters(), 0.01f);
-
-    // one forward/backward pass to populate gradients
-    auto pred = (*net)(x);
-    auto l = (*loss)(y, pred);
-    l->backward();
-
-    // verify gradients are non-zero before zeroing
-    bool anyNonZero = false;
-    for(auto& p : net->parameters()) {
-        if(p->getGrads()) {
-            for(tensorSize_t i = 0; i < p->getGrads()->getSize(); i++) {
-                if((*p->getGrads())[i] != 0.0f) {
-                    anyNonZero = true;
-                    break;
-                }
-            }
-        }
-    }
-    EXPECT_TRUE(anyNonZero) << "Expected some non-zero gradients before zeroGrad";
-
-    // zero gradients
-    optim->zeroGrad();
-
-    // verify all gradients are zero after zeroing
-    for(auto& p : net->parameters()) {
-        if(p->getGrads()) {
-            for(tensorSize_t i = 0; i < p->getGrads()->getSize(); i++) {
-                EXPECT_FLOAT_EQ((*p->getGrads())[i], 0.0f)
-                    << "Gradient not zeroed at index " << i;
-            }
-        }
-    }
-}
\ No newline at end of file
diff --git a/tests/python/__init__.py b/tests/python/__init__.py
deleted file mode 100644
index e69de29..0000000
diff --git a/tests/python/cuda/test_autograd_cuda.py b/tests/python/cuda/test_autograd_cuda.py
new file mode 100644
index 0000000..1fa6c51
--- /dev/null
+++ b/tests/python/cuda/test_autograd_cuda.py
@@ -0,0 +1,111 @@
+"""
+Robert Baumgartner, r.baumgartner-1@tudelft.nl
+"""
+
+import sys
+from pathlib import Path
+
+sys.path.insert(0, str(Path(__file__).parent.parent.parent.parent / "python_lib"))
+
+from dl_lib import Tensor, Device
+import pytest
+
+
+class TestAutogradCuda:
+  def test_backward(self):
+    t = Tensor([2, 2], Device.CUDA, True)
+
+    loss = (t * 2).sum()
+    loss.backward()
+
+    assert t.grads is not None
+
+  def test_nograd_throws(self):
+    t1 = Tensor([1], [3.0], Device.CUDA, False)
+    t2 = Tensor([1], [3.0], Device.CUDA, False)
+
+    t3 = t1 * t2
+
+    assert not t3.requiresGrad
+    with pytest.raises(RuntimeError):
+      t3.backward()
+
+  def test_add(self):
+    t1 = Tensor([1], [3.0], Device.CUDA, True)
+    t2 = Tensor([1], [2.0], Device.CUDA, True)
+
+    t3 = t1 + t2
+    loss = t3 * t3
+
+    loss.backward()
+
+    assert t1.grads.getitem(0) == pytest.approx(10.0, abs=1e-4)
+    assert t2.grads.getitem(0) == pytest.approx(10.0, abs=1e-4)
+
+  def test_scalar_mul(self):
+    t1 = Tensor([1], [2.0], Device.CUDA, True)
+    t2 = Tensor([1], [3.0], Device.CUDA, True)
+
+    t3 = t1 * t2
+    loss = t3 * t3
+
+    loss.backward()
+
+    assert t1.grads.getitem(0) == pytest.approx(36.0, abs=1e-4)
+    assert t2.grads.getitem(0) == pytest.approx(24.0, abs=1e-4)
+
+  def test_matmul(self):
+    t1 = Tensor([2, 3], [1, 2, 3, 4, 5, 6], Device.CUDA, True)
+    t2 = Tensor([3, 2], [1, 2, 3, 4, 5, 6], Device.CUDA, True)
+
+    t3 = t1 @ t2
+    loss = t3.sum()
+
+    loss.backward()
+
+    # dL/dt1 = dloss/dt3 @ t2^T = Ones({2, 2}) @ t2^T
+    assert t1.grads.getitem([0, 0]) == pytest.approx(3.0, abs=1e-4)
+    assert t1.grads.getitem([0, 1]) == pytest.approx(7.0, abs=1e-4)
+    assert t1.grads.getitem([0, 2]) == pytest.approx(11.0, abs=1e-4)
+    assert t1.grads.getitem([1, 0]) == pytest.approx(3.0, abs=1e-4)
+    assert t1.grads.getitem([1, 1]) == pytest.approx(7.0, abs=1e-4)
+    assert t1.grads.getitem([1, 2]) == pytest.approx(11.0, abs=1e-4)
+
+    # dL/dt2 = t1^T @ dloss/dt3 = t1^T @ Ones({2, 2})
+    assert t2.grads.getitem([0, 0]) == pytest.approx(5.0, abs=1e-4)
+    assert t2.grads.getitem([0, 1]) == pytest.approx(5.0, abs=1e-4)
+    assert t2.grads.getitem([1, 0]) == pytest.approx(7.0, abs=1e-4)
+    assert t2.grads.getitem([1, 1]) == pytest.approx(7.0, abs=1e-4)
+    assert t2.grads.getitem([2, 0]) == pytest.approx(9.0, abs=1e-4)
+    assert t2.grads.getitem([2, 1]) == pytest.approx(9.0, abs=1e-4)
+
+  def test_chainrule(self):
+    x = Tensor([1], [2.0], Device.CUDA, True)
+
+    y = x * x       # y = x^2
+    z = x + y       # z = x^2 + x
+    loss = z * z    # loss = (x^2 + x)^2
+
+    loss.backward()
+
+    # dloss/dx = 2(x^2 + x)(2x + 1) = 60 at x=2
+    assert x.grads.getitem(0) == pytest.approx(60.0, abs=1e-4)
+
+  def test_multivariate_chainrule(self):
+    x = Tensor([2], [1.0, 2.0], Device.CUDA, True)
+    y = x * 3
+
+    loss = Tensor([1], [0.0], Device.CUDA, True)
+    for i in range(len(y)):
+      loss = loss + y[i]
+    loss.backward()
+
+    assert x.grads.getitem(0) == pytest.approx(3.0, abs=1e-4)
+    assert x.grads.getitem(1) == pytest.approx(3.0, abs=1e-4)
+
+    assert y.grads.getitem(0) == pytest.approx(1.0, abs=1e-4)
+    assert y.grads.getitem(1) == pytest.approx(1.0, abs=1e-4)
+
+
+if __name__ == '__main__':
+  raise RuntimeError("Not a standalone script")
diff --git a/tests/python/cuda/test_tensorops_cuda.py b/tests/python/cuda/test_tensorops_cuda.py
new file mode 100644
index 0000000..44503da
--- /dev/null
+++ b/tests/python/cuda/test_tensorops_cuda.py
@@ -0,0 +1,41 @@
+"""
+Robert Baumgartner, r.baumgartner-1@tudelft.nl
+"""
+
+import sys
+from pathlib import Path
+
+sys.path.insert(0, str(Path(__file__).parent.parent.parent.parent / "python_lib"))
+
+from dl_lib import Tensor, Device
+
+
+class TestTensorOpsCuda:
+  def test_ones(self):
+    t = Tensor.ones([2, 2], Device.CUDA)
+    assert t.dims == [2, 2]
+    assert t.device == Device.CUDA
+
+  def test_ctor(self):
+    t = Tensor([2], [1.0, 2.0], Device.CUDA, False)
+
+    assert t.getitem(0) == 1.0
+    assert t.getitem(1) == 2.0
+
+    assert t.device == Device.CUDA
+    assert t.requiresGrad == False
+
+  def test_multiplication(self):
+    a = Tensor.ones([2, 2], Device.CUDA) * 3
+    b = Tensor.ones([2, 2], Device.CUDA) * 0.5
+    c = a * b
+
+    assert c.dims == [2, 2]
+    assert c.getitem([0, 0]) == 1.5
+    assert c.getitem([0, 1]) == 1.5
+    assert c.getitem([1, 0]) == 1.5
+    assert c.getitem([1, 1]) == 1.5
+
+
+if __name__ == '__main__':
+  raise RuntimeError("Not a standalone script")
diff --git a/tests/python/cuda/test_training_cuda.py b/tests/python/cuda/test_training_cuda.py
new file mode 100644
index 0000000..bb48d92
--- /dev/null
+++ b/tests/python/cuda/test_training_cuda.py
@@ -0,0 +1,149 @@
+"""
+Robert Baumgartner, r.baumgartner-1@tudelft.nl
+"""
+
+import sys
+from pathlib import Path
+
+sys.path.insert(0, str(Path(__file__).parent.parent.parent.parent / "python_lib"))
+
+from dl_lib import Tensor, Device
+from dl_lib.nn import FfLayer, Sequential
+from dl_lib.nn.activation import LeakyReLU
+from dl_lib.train.loss import BceWithSigmoid, CrossEntropyWithSoftmax
+from dl_lib.train.optim import SGD, RmsProp
+
+from dl_lib.sys import setSeed
+import pytest
+
+setSeed(42)
+
+def train(net, loss_fn, optim, x, y, epochs):
+  for _ in range(epochs):
+    ypred = net.forward(x)
+    loss = loss_fn(y, ypred)
+
+    loss.backward()
+    optim.step()
+    optim.zeroGrad()
+
+  return loss
+
+
+def make_binary_net():
+  net = Sequential()
+  net.append(FfLayer(2, 4, Device.CUDA))
+  net.append(LeakyReLU(0.01))
+  net.append(FfLayer(4, 1, Device.CUDA))
+  return net
+
+
+def make_multiclass_net():
+  net = Sequential()
+  net.append(FfLayer(2, 8, Device.CUDA))
+  net.append(LeakyReLU(0.01))
+  net.append(FfLayer(8, 3, Device.CUDA))
+  return net
+
+
+def make_xor_data():
+  x = Tensor([4, 2], [0.0, 0.0,
+                      0.0, 1.0,
+                      1.0, 0.0,
+                      1.0, 1.0], Device.CUDA, False)
+  y = Tensor([4, 1], [0.0,
+                      1.0,
+                      1.0,
+                      0.0], Device.CUDA, False)
+  return x, y
+
+
+def make_multiclass_data():
+  x = Tensor([6, 2], [1.0, 0.0,
+                      1.0, 0.1,
+                      0.0, 1.0,
+                      0.1, 1.0,
+                      0.5, 0.5,
+                      0.4, 0.6], Device.CUDA, False)
+  y = Tensor([6, 3], [1.0, 0.0, 0.0,
+                      1.0, 0.0, 0.0,
+                      0.0, 1.0, 0.0,
+                      0.0, 1.0, 0.0,
+                      0.0, 0.0, 1.0,
+                      0.0, 0.0, 1.0], Device.CUDA, False)
+  return x, y
+
+
+class TestOverfitBinaryCuda:
+  def test_binary_sgd_overfits(self):
+    x, y = make_xor_data()
+    net = make_binary_net()
+    loss_fn = BceWithSigmoid()
+    optim = SGD(net.parameters(), 0.05)
+
+    final_loss = train(net, loss_fn, optim, x, y, epochs=3000)
+
+    assert final_loss.getitem(0) < 0.05, \
+      f"SGD failed to overfit XOR on CUDA, loss={final_loss.getitem(0)}"
+
+  def test_binary_rmsprop_overfits(self):
+    x, y = make_xor_data()
+    net = make_binary_net()
+    loss_fn = BceWithSigmoid()
+    optim = RmsProp(net.parameters(), 0.001, 0.95)
+
+    final_loss = train(net, loss_fn, optim, x, y, epochs=3000)
+
+    assert final_loss.getitem(0) < 0.05, \
+      f"RmsProp failed to overfit XOR on CUDA, loss={final_loss.getitem(0)}"
+
+  def test_multiclass_rmsprop_overfits(self):
+    x, y = make_multiclass_data()
+    net = make_multiclass_net()
+    loss_fn = CrossEntropyWithSoftmax()
+    optim = RmsProp(net.parameters(), 0.001, 0.95)
+
+    final_loss = train(net, loss_fn, optim, x, y, epochs=3000)
+
+    assert final_loss.getitem(0) < 0.05, \
+      f"RmsProp failed to overfit multiclass on CUDA, loss={final_loss.getitem(0)}"
+
+  def test_loss_decreases(self):
+    x, y = make_xor_data()
+    net = make_binary_net()
+    loss_fn = BceWithSigmoid()
+    optim = SGD(net.parameters(), 0.001)
+
+    initial_pred = net.forward(x)
+    initial_loss = loss_fn(y, initial_pred).getitem(0)
+    train(net, loss_fn, optim, x, y, epochs=2000)
+
+    final_pred = net.forward(x)
+    final_loss = loss_fn(y, final_pred).getitem(0)
+
+    assert final_loss < initial_loss, \
+      f"Loss did not decrease on CUDA: {initial_loss} -> {final_loss}"
+
+  def test_gradients_are_zeroed_between_steps(self):
+    x, y = make_xor_data()
+    net = make_binary_net()
+    loss_fn = BceWithSigmoid()
+    optim = SGD(net.parameters(), 0.01)
+
+    for _ in range(2):
+      pred = net.forward(x)
+      loss = loss_fn(y, pred)
+
+      loss.backward()
+      optim.step()
+      optim.zeroGrad()
+
+    for p in net.parameters():
+      if p.grads is not None:
+        for i in range(p.grads.size):
+          assert p.grads.getitem(i) == pytest.approx(0.0, abs=1e-5), \
+            f"Gradient not zeroed at index {i}"
+
+
+if __name__ == '__main__':
+  raise RuntimeError("Not a standalone script")
diff --git a/tests/python/test_training.py b/tests/python/test_training.py
index 822c386..bf4bb97 100644
--- a/tests/python/test_training.py
+++ b/tests/python/test_training.py
@@ -77,7 +77,7 @@ def test_binary_sgd_overfits(self):
     loss_fn = BceWithSigmoid()
     optim = SGD(net.parameters(), 0.05)
 
-    final_loss = train(net, loss_fn, optim, x, y, epochs=2000)
+    final_loss = train(net, loss_fn, optim, x, y, epochs=3000)
 
     assert final_loss.getitem(0) < 0.05, \
       f"SGD failed to overfit XOR, loss={final_loss.getitem(0)}"
@@ -86,9 +86,9 @@ def test_binary_rmsprop_overfits(self):
     x, y = make_xor_data()
     net = make_binary_net()
     loss_fn = BceWithSigmoid()
-    optim = RmsProp(net.parameters(), 0.0001, 0.95)
+    optim = RmsProp(net.parameters(), 0.001, 0.95)
 
-    final_loss = train(net, loss_fn, optim, x, y, epochs=5000)
+    final_loss = train(net, loss_fn, optim, x, y, epochs=3000)
 
     assert final_loss.getitem(0) < 0.05, \
       f"RmsProp failed to overfit XOR, loss={final_loss.getitem(0)}"
@@ -97,9 +97,9 @@ def test_multiclass_rmsprop_overfits(self):
     x, y = make_multiclass_data()
     net = make_multiclass_net()
     loss_fn = CrossEntropyWithSoftmax()
-    optim = RmsProp(net.parameters(), 0.0003, 0.95)
+    optim = RmsProp(net.parameters(), 0.001, 0.95)
 
-    final_loss = train(net, loss_fn, optim, x, y, epochs=10000)
+    final_loss = train(net, loss_fn, optim, x, y, epochs=3000)
 
     assert final_loss.getitem(0) < 0.05, \
         f"RmsProp failed to overfit multiclass, loss={final_loss.getitem(0)}"