From cf0d4fdf3bf8fa6e1d6abf631042de28176923c3 Mon Sep 17 00:00:00 2001 From: Masahiro Masuda Date: Fri, 25 Dec 2020 10:12:01 +0900 Subject: [PATCH] get valid count test working --- python/tvm/topi/cuda/nms.py | 2 +- python/tvm/topi/cuda/scan.py | 3 +-- 2 files changed, 2 insertions(+), 3 deletions(-) diff --git a/python/tvm/topi/cuda/nms.py b/python/tvm/topi/cuda/nms.py index 4b6452918a4a..2d90bc485625 100644 --- a/python/tvm/topi/cuda/nms.py +++ b/python/tvm/topi/cuda/nms.py @@ -22,7 +22,7 @@ from tvm.tir import if_then_else from .sort import argsort, argsort_thrust -from .prefix_scan import exclusive_scan +from .scan import exclusive_scan from ..utils import ceil_div diff --git a/python/tvm/topi/cuda/scan.py b/python/tvm/topi/cuda/scan.py index 566aad4d9957..f3dfeb875743 100644 --- a/python/tvm/topi/cuda/scan.py +++ b/python/tvm/topi/cuda/scan.py @@ -237,9 +237,8 @@ def exclusive_scan(data, axis=-1, return_reduction=False): [data.shape, (data.shape[0],)], [data], lambda ins, outs: exclusive_sum_scan2d_ir(ins[0], outs[0], outs[1]), - dtype=[data.dtype], + dtype=[data.dtype, data.dtype], in_buffers=[data_buf], - out_buffers=[output_buf], name="exclusive_scan", tag="exclusive_scan_gpu", )