From 08b504ebe7c3bb431645a0911aeeccab0cfed9c5 Mon Sep 17 00:00:00 2001 From: Altan Haan Date: Tue, 29 Jun 2021 11:55:48 -0700 Subject: [PATCH] silence terrible pylint suggestion --- python/tvm/topi/nn/softmax.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/python/tvm/topi/nn/softmax.py b/python/tvm/topi/nn/softmax.py index 210235a899bf..a13b17686708 100644 --- 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))