Fix test_softmax_correctness_with_axis_tuple test case failed on torc…

…h gpu ci. (keras-team#20025)
SamanehSaadat · Jul 22, 2024 · f9501f5 · f9501f5
1 parent 93786e3
commit f9501f5
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/keras/src/ops/nn_test.py b/keras/src/ops/nn_test.py
@@ -1255,7 +1255,9 @@ def test_softmax_correctness_with_axis_tuple(self):
  combination = combinations(range(3), 2)
  for axis in list(combination):
  result = keras.ops.nn.softmax(input, axis=axis)
- normalized_sum_by_axis = np.sum(np.asarray(result), axis=axis)
+ normalized_sum_by_axis = np.sum(
+ ops.convert_to_numpy(result), axis=axis
+ )
  self.assertAllClose(normalized_sum_by_axis, 1.0)
 
  def test_log_softmax(self):