GPU上的PyTorch列表切片速度比CPU上的慢

Question

我想优化NVIDIA Jetson Xavier NX（开发工具包）上的ML代码（SSD）。瓶颈之一似乎是GPU设备上PyTorch（1.6.0）张量的列表切片

同样的问题也发生在NVIDIA GeForce GTX 1050 Ti（GP107）上，CPU速度提高了约2倍

让我先创建变量

import torch
from time import time

cuda0 = torch.device('cuda:0')

probs = torch.ones([3000], dtype=torch.float64, device=cuda0)
mask = torch.ones([3000], dtype=torch.bool, device=cuda0)

probs_cpu = probs.cpu()
mask_cpu = mask.cpu()

然后运行逻辑（每次运行的结果大致相同）

before = time()
probs[mask]
print(f'GPU {time() - before:.5f}') # output: GPU 0.00263


before = time()
probs_cpu[mask_cpu]
print(f'CPU {time() - before:.5f}') # output: CPU 0.00066

根据上述代码，为什么GPU上的列表切片速度比NVIDIA Jetson Xavier NX Developer kit上使用Pyrotch library vesrion 1.6.0的CPU慢4倍？如何加速

代码详细信息：参见predictor.py中的第51行，它是SSD Implementation in PyTorch的一部分

在CPU上运行它？：如果在CPU上运行它，整个算法不会更快，因为从GPU下载需要太长时间（~0.00805秒）