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

Question

我想在NVIDIA Jetson Xavier NX（开发套件）上优化ML代码（PyTorch中的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

与上面根据代码在NVIDIA Jetson Xavier NX Developer套件上使用PyTorch库vesrion 1.6.0的CPU相比，为什么列表在GPU上的切片速度要慢4倍？如何加快速度？

代码详细信息：请参阅predictor.py的SSD Implementation in PyTorch中的第51行

在CPU上运行它？：如果我在CPU上运行它，整个算法将不会更快，因为从GPU上下载的时间太长（〜0.00805 s）。