关于使用libtorch实现Yolov5的NMS算法比pytorch慢的问题

Question

python代码片段，红色部分的切片操作：
def non_max_suppression(prediction, conf_thres=0.25, iou_thres=0.45,         classes=None, agnostic=False, multi_label=False,
                        labels=(), max_det=300):
    
    nc = prediction.shape[2] - 5  # number of classes
    xc = prediction[..., 4] > conf_thres  # candidates
    
    output = [torch.zeros((0, 6), device=prediction.device)] * prediction.shape[0]
    print('-----------------------------------------')
    for xi, x in enumerate(prediction):  # image index, image inference
        
        begin_time = datetime.datetime.now()
        x = x[xc[xi]]  # confidence
        end_time = datetime.datetime.now()
        d_time = end_time - begin_time
        print(str(xi) + ': ' + str(d_time.microseconds))
其中prediction的形状是torch.Size([16, 15120, 85])
 x = x[xc[xi]]这一句在循环中的执行时间如下，单位是微秒：

C++ Libtorch代码片段：

其中torch::masked_select一行对应于Pytorch版本中的x = x[xc[xi]]，执行的时间如下，单位是毫秒：

第一次切片操作Libtorch处理时间是Python版本的一倍还多。

第一次forward的时间也比pytorch长的多，请教有没有人解决过这个问题？
我现在使用的libtorch相关的lib和dll都是从python那边拷贝过来的。