search

grimoire / torch2trt_dynamic

A pytorch to tensorrt convert with dynamic shape support
MIT License
254 stars 34 forks source link

topk issue #8

Closed niaoyu closed 3 years ago

niaoyu commented 3 years ago

Hi, I used these repo to transfer the following code from torch to trt.

    topk_scores, topk_inds = torch.topk(scores.view(batch, cat, height*width), K)
    topk_inds = topk_inds.int()
    topk_inds = topk_inds % (height * width)

However, the error shown here. [TensorRT] ERROR: (Unnamed Layer* 98) [TopK]: reduceAxes out of range

Are there some issues in convert/topk.py?

grimoire commented 3 years ago

Hi The code works fine on my side. Could you please provide more detail? Such as shape of score, value of K etc.

niaoyu commented 3 years ago

Hi The code works fine on my side. Could you please provide more detail? Such as shape of score, value of K etc.

Sry, it is caused by my mistake! Such a great work, many thanks!