Adam Maroni
--- a/unit_tests/operator/Test_SoftmaxImpl.cpp
+++ b/unit_tests/operator/Test_SoftmaxImpl.cpp
+}
+
+
+
+TEST_CASE("[cpu/operator] Softmax(backward)", "[Softmax][CPU]") {
+    SECTION("1D Tensor") {
+      std::shared_ptr<Softmax_Op> op = std::make_shared<Softmax_Op>(0);
+      op->setDataType(DataType::Float32);
+      op->setBackend("cpu");
+
+      std::shared_ptr<Tensor> softMaxForwardInputTensor = 
+        std::make_shared<Tensor>(Array1D<float,3> { {3.0, 1.0, 0.2} });
--- a/include/aidge/backend/cpu/operator/SoftmaxImpl_kernels.hpp
+++ b/include/aidge/backend/cpu/operator/SoftmaxImpl_kernels.hpp
+{
+    const O* softmaxOut = static_cast<const O*>(softmaxOut_); 
+    const O* target = static_cast<const O*>(target_); 
+    I* dL = static_cast<I*>(gradientLoss_);
+
+    // Compute the number of elements after the softmax axis (post-axis size)
+    std::size_t postAxisElems = 1;
+    for (std::size_t i = axisIdx + 1; i < inputDims.size(); ++i) {
+        postAxisElems *= inputDims[i];
+    }
+
+    // Compute the number of elements after the softmax axis (pre-axis size)
+    std::size_t preAxisElems = 1;
+    for (std::size_t i = 0; i < axisIdx; ++i) {
+        preAxisElems *= inputDims[i];
+    }