diff --git a/src/matrix.cu b/src/matrix.cu
index 4e73b74..7a536ae 100644
--- a/src/matrix.cu
+++ b/src/matrix.cu
@@ -87,17 +87,17 @@ static __device__ inline float fastexp(float x) {
 }
 
 __device__ void softmax(float* a, int rows) {
-    float res = (float)0;
-    for (int i = 0; i < rows; i++) {
-        res += exp(a[i]);
+    float sum = 0.0;
+    for (size_t i = 0; i < rows; i++) {
+        sum += __expf(a[i]);
     }
-    for (int i = 0; i < rows; i++) {
-        a[i] /= res;
+    for (size_t i = 0; i < rows; i++) {
+        a[i] = __expf(a[i] - __logf(sum));
     }
 }
 
 __device__ int argmax(float* a, int rows) {
-    int res = a[0];
+    float res = a[0];
     int idx = 0;
     for (int i = 0; i < rows; i++) {
         if (res < a[i]) {