diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml
index 52cdd92a..5f6a47c8 100644
--- a/.github/workflows/ci.yaml
+++ b/.github/workflows/ci.yaml
@@ -30,7 +30,7 @@ jobs:
           pytest --junitxml=pytest.xml --cov-report=term-missing:skip-covered --cov=choice_learn tests/ | tee pytest-coverage.txt
 
       - name: Pytest coverage comment
-        uses: MishaKav/pytest-coverage-comment@main
+        uses: VincentAuriau/pytest-coverage-comment@main
         with:
           pytest-coverage-path: ./pytest-coverage.txt
           junitxml-path: ./pytest.xml
diff --git a/.github/workflows/pr_ci.yaml b/.github/workflows/pr_ci.yaml
index 61ffeb5e..d73349e9 100644
--- a/.github/workflows/pr_ci.yaml
+++ b/.github/workflows/pr_ci.yaml
@@ -37,3 +37,4 @@ jobs:
           report-only-changed-files: true
           unique-id-for-comment: ${{ matrix.python-version }}
           title: 'Coverage Report for Python ${{ matrix.python-version }}'
+
diff --git a/choice_learn/tf_ops.py b/choice_learn/tf_ops.py
index de39ac1d..44e626f9 100644
--- a/choice_learn/tf_ops.py
+++ b/choice_learn/tf_ops.py
@@ -33,16 +33,17 @@ def softmax_with_availabilities(
         Probabilities of each product for each choice computed from Logits
     """
     # Substract max utility to avoid overflow
-    numerator = tf.exp(
-        items_logit_by_choice - tf.reduce_max(items_logit_by_choice, axis=axis, keepdims=True)
-    )
+    normalizer = tf.reduce_max(items_logit_by_choice, axis=axis, keepdims=True)
+    numerator = tf.exp(items_logit_by_choice - normalizer)
+
     # Set unavailable products utility to 0
     numerator = tf.multiply(numerator, available_items_by_choice)
     # Sum of total available utilities
     denominator = tf.reduce_sum(numerator, axis=axis, keepdims=True)
+
     # Add 1 to the denominator to take into account the exit choice
     if normalize_exit:
-        denominator += 1
+        denominator += tf.exp(-normalizer)
     # Avoir division by 0 when only unavailable items have highest utilities
     elif eps:
         denominator += eps
diff --git a/tests/unit_tests/test_tf_ops.py b/tests/unit_tests/test_tf_ops.py
index 9775fbaa..d4de5e24 100644
--- a/tests/unit_tests/test_tf_ops.py
+++ b/tests/unit_tests/test_tf_ops.py
@@ -11,29 +11,30 @@ def test_softmax():
         [[1, 1, 1, 1], [2, 1, 2, 1], [np.log(0.125), np.log(0.125), np.log(0.25), np.log(0.5)]]
     )
     availabilities = np.array([[1.0, 1.0, 1.0, 1.0], [1.0, 0.0, 1.0, 0.0], [1.0, 1.0, 1.0, 1.0]])
-    """
-    probabilities = np.array([[0.25, 0.25, 0.25, 0.25],
-                              [0.5, 0.0, 0.5, 0.0],
-                              [0.125, 0.125, 0.250, 0.5]])
-    """
+
+    probabilities = np.array(
+        [[0.25, 0.25, 0.25, 0.25], [0.5, 0.0, 0.5, 0.0], [0.125, 0.125, 0.250, 0.5]]
+    )
+
     softmax_probabilities = softmax_with_availabilities(
         items_logit_by_choice=logits, available_items_by_choice=availabilities
     ).numpy()
-    assert (softmax_probabilities[0] < 0.26).all()
-    assert (softmax_probabilities[0] > 0.24).all()
-
-    assert (softmax_probabilities[1][[0, 2]] < 0.51).all()
-    assert (softmax_probabilities[1][[0, 2]] > 0.49).all()
-    assert (softmax_probabilities[1][[1, 3]] < 0.01).all()
-
-    assert (softmax_probabilities[2][0] < 0.126).all()
-    assert (softmax_probabilities[2][0] > 0.124).all()
-    assert (softmax_probabilities[2][1] < 0.126).all()
-    assert (softmax_probabilities[2][1] > 0.124).all()
-    assert (softmax_probabilities[2][2] < 0.251).all()
-    assert (softmax_probabilities[2][2] > 0.249).all()
-    assert (softmax_probabilities[2][3] < 0.501).all()
-    assert (softmax_probabilities[2][3] > 0.499).all()
+    assert (np.abs(softmax_probabilities - probabilities) < 0.01).all()
+
+
+def test_softmax_exit():
+    """Test the softmax function with normalized exit."""
+    logits = np.array(
+        [[np.log(1.0), np.log(1.0), np.log(1.0), np.log(1.0)], [np.log(2.0), 1.0, np.log(2.0), 1.0]]
+    )
+    availabilities = np.array([[1.0, 1.0, 1.0, 1.0], [1.0, 0.0, 1.0, 0.0]])
+
+    probabilities = np.array([[0.20, 0.20, 0.20, 0.20], [0.4, 0.0, 0.4, 0.0]])
+
+    softmax_probabilities = softmax_with_availabilities(
+        items_logit_by_choice=logits, available_items_by_choice=availabilities, normalize_exit=True
+    ).numpy()
+    assert (np.abs(softmax_probabilities - probabilities) < 0.01).all()
 
 
 def test_custom_categorical_crossentropy():