Skip to content

Commit

Permalink
space line added
Browse files Browse the repository at this point in the history
  • Loading branch information
Laurawly committed Mar 12, 2019
1 parent 137a382 commit a81aed6
Showing 1 changed file with 4 additions and 0 deletions.
4 changes: 4 additions & 0 deletions topi/python/topi/cuda/nms.py
Original file line number Diff line number Diff line change
Expand Up @@ -73,6 +73,7 @@ def get_valid_counts_pre(data, flag, idx, score_threshold):

return ib.get()


def get_valid_counts_ir(data, flag, idx, valid_count, out):
"""Low level IR to get valid count of bounding boxes
given a score threshold. Also moves valid boxes to the
Expand Down Expand Up @@ -133,6 +134,8 @@ def get_valid_counts_ir(data, flag, idx, valid_count, out):
valid_count[i] = idx[i * num_anchors + num_anchors - 1]

return ib.get()


@get_valid_counts.register(["cuda", "gpu"])
def get_valid_counts_gpu(data, score_threshold=0):
"""Get valid count of bounding boxes given a score threshold.
Expand Down Expand Up @@ -180,6 +183,7 @@ def get_valid_counts_gpu(data, score_threshold=0):

return [valid_count, out_tensor]


def sort_ir(data, index, output):
"""Low level IR to do sorting on the GPU, same usage as tvm.contrib.sort.argsort on the CPU.
Expand Down

0 comments on commit a81aed6

Please sign in to comment.