silence terrible pylint suggestion

apache · Jun 30, 2021 · 08b504e · 08b504e
1 parent 8d8fdc1
commit 08b504e
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/python/tvm/topi/nn/softmax.py b/python/tvm/topi/nn/softmax.py
@@ -136,9 +136,9 @@ def log_softmax(x, axis=-1):
     output : tvm.te.Tensor
         2-D output with same shape
     """
-
     assert len(x.shape) == 2, "only support 2-dim log softmax"
-    assert axis == 1 or axis == len(x.shape) - 1, "only support last axis log softmax"
+    # pylint: disable=R1714
+    assert axis == -1 or axis == len(x.shape) - 1, "only support last axis log softmax"
     m, n = x.shape
     k = te.reduce_axis((0, n), name="k")
     max_elem = te.compute((m,), lambda i: tvm.te.max(x[i, k], axis=k))