diff --git a/python/tvm/relay/frontend/pytorch.py b/python/tvm/relay/frontend/pytorch.py index 0213dcc488fd..b9650e6e9a9c 100644 --- a/python/tvm/relay/frontend/pytorch.py +++ b/python/tvm/relay/frontend/pytorch.py @@ -2680,7 +2680,7 @@ def nonzero(self, inputs, input_types, is_numpy_style=False): return ret def nonzero_numpy(self, inputs, input_types): - return self.nonzero(inputs, input_types, is_numpy_style=False) + return self.nonzero(inputs, input_types, is_numpy_style=True) def scatter(self, inputs, input_types): assert len(inputs) == 4 or len(inputs) == 5, ( diff --git a/tests/python/frontend/pytorch/test_forward.py b/tests/python/frontend/pytorch/test_forward.py index 6178a58b6d13..9bf40cfcdd85 100644 --- a/tests/python/frontend/pytorch/test_forward.py +++ b/tests/python/frontend/pytorch/test_forward.py @@ -4445,6 +4445,7 @@ def forward(self, data): inp = torch.Tensor(np.array([[0, 1, 0], [2, 0, 9], [-1, -1, 0]]).astype("float32")) verify_trace_model(Nonzero(), [inp], ["llvm"]) + verify_trace_model(Nonzero(as_tuple=True), [inp], ["llvm"]) def test_forward_scatter():