diff --git a/src/kernels/activation_functions.cu b/src/kernels/activation_functions.cu
index 296c7cb..7b5956b 100644
--- a/src/kernels/activation_functions.cu
+++ b/src/kernels/activation_functions.cu
@@ -1,4 +1,4 @@
-#include <functional>
+#include <cmath>
 
 #include "activation_functions.cuh"
 #include "cuda_helper.cuh"
@@ -38,7 +38,7 @@ __global__ void CUDANet::Kernels::softmax_exp(
     int tid    = blockDim.x * blockIdx.x + threadIdx.x;
 
     for (int i = tid; i < len; i += stride) {
-        dst[i] = exp(src[i]);
+        dst[i] = std::exp(src[i]);
     }
 }
 
diff --git a/src/layers/activation.cu b/src/layers/activation.cu
index 94e9cce..7ceb1b8 100644
--- a/src/layers/activation.cu
+++ b/src/layers/activation.cu
@@ -41,7 +41,7 @@ void Layers::Activation::activate(float* __restrict__ d_input) {
                 d_input, d_input, length
             );
 
-            Kernels::softmax_sum<<<gridSize, BLOCK_SIZE>>>(
+            Kernels::softmax_sum<<<gridSize / 2, BLOCK_SIZE>>>(
                 d_input, d_softmax_sum, length
             );
 
diff --git a/test/CMakeLists.txt b/test/CMakeLists.txt
index af1ab20..4db8714 100644
--- a/test/CMakeLists.txt
+++ b/test/CMakeLists.txt
@@ -6,6 +6,7 @@ add_executable(test_main
     layers/test_dense.cu
     layers/test_conv2d.cu
     layers/test_input.cu
+    layers/test_activation.cu
     kernels/test_activation_functions.cu
     kernels/test_padding.cu
     kernels/test_matmul.cu
diff --git a/test/layers/test_activation.cu b/test/layers/test_activation.cu
new file mode 100644
index 0000000..e501799
--- /dev/null
+++ b/test/layers/test_activation.cu
@@ -0,0 +1,34 @@
+#include "activation.cuh"
+#include <gtest/gtest.h>
+#include <cuda_runtime.h>
+#include <vector>
+
+TEST(ActivationTest, SoftmaxTest) {
+    CUDANet::Layers::Activation activation(
+        CUDANet::Layers::ActivationType::SOFTMAX, 5
+    );
+
+    std::vector<float> input = {0.573f, 0.619f, 0.732f, 0.055f, 0.243f};
+
+    float* d_input;
+    cudaMalloc((void**)&d_input, sizeof(float) * 5);
+    cudaMemcpy(d_input, input.data(), sizeof(float) * 5, cudaMemcpyHostToDevice);
+
+    activation.activate(d_input);
+    std::vector<float> output(5);
+    cudaMemcpy(
+        output.data(), d_input, sizeof(float) * 5, cudaMemcpyDeviceToHost
+    );
+    
+    float sum = 0.0f;
+
+    std::vector<float> expected = {0.22055f, 0.23094f, 0.25856f, 0.13139f, 0.15856f};
+    for (int i = 0; i < 5; ++i) {
+        sum += output[i];
+        EXPECT_NEAR(output[i], expected[i], 1e-5f);
+    }
+
+    EXPECT_NEAR(sum, 1.0f, 1e-5f);
+
+    cudaFree(d_input);
+}
\ No newline at end of file
diff --git a/tools/generate_conv2d_test.py b/tools/generate_conv2d_test.py
index 1b3e6d2..fc4770f 100644
--- a/tools/generate_conv2d_test.py
+++ b/tools/generate_conv2d_test.py
@@ -110,10 +110,20 @@ def gen_strided_test_result():
     output = conv2d(in_channels, out_channels, kernel_size, stride, padding, input, weights)
     print_cpp_vector(output)
 
+def gen_softmax_test_result():
+    input = torch.tensor([
+        0.573, 0.619, 0.732, 0.055, 0.243
+    ])
+
+    output = torch.nn.Softmax(dim=0)(input)
+    print_cpp_vector(output)
+
 
 if __name__ == "__main__":
     print("Generating test results...")
     print("Padded convolution test:")
     gen_padded_test_result()
     print("Strided convolution test:")
-    gen_strided_test_result()
\ No newline at end of file
+    gen_strided_test_result()
+    print("Softmax test:")
+    gen_softmax_test_result()
\ No newline at end of file