diff --git a/tensorflow_addons/layers/python/sparsemax_test.py b/tensorflow_addons/layers/python/sparsemax_test.py
index 3cb375418e..6796982c2f 100644
--- a/tensorflow_addons/layers/python/sparsemax_test.py
+++ b/tensorflow_addons/layers/python/sparsemax_test.py
@@ -60,7 +60,7 @@ def test_sparsemax_layer_against_numpy(self, dtype=None):
         z = random.uniform(low=-3, high=3, size=(test_obs, 10)).astype(dtype)
 
         test_utils.layer_test(
-            layer_cls=Sparsemax,
+            Sparsemax,
             input_data=z,
             expected_output=_np_sparsemax(z).astype(dtype))