get valid count test working

trevor-m · Dec 25, 2020 · cf0d4fd · cf0d4fd
1 parent eb142d3
commit cf0d4fd
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 3 deletions.
diff --git 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
@@ -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",
             )