Нарезка списка PyTorch на GPU медленнее, чем на CPU

Я хотел бы оптимизировать код машинного обучения (SSD в PyTorch) на NVIDIA Jetson Xavier NX (комплект для разработки ). Одним из узких мест, похоже, является нарезка списков на тензорах PyTorch (1.6.0) на устройстве с графическим процессором.

Та же проблема возникла на NVIDIA GeForce GTX 1050 Ti (GP107), процессор был примерно в 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

Почему нарезка списка в ~ 4 раза медленнее на GPU по сравнению с CPU с использованием библиотеки PyTorch vesrion 1.6.0 в комплекте разработчика NVIDIA Jetson Xavier NX в соответствии с приведенным выше кодом? Как это ускорить?

Сведения о коде: см. строку 51 в predictor.py, который является частью реализации SSD в PyTorch

Запустить его на ЦП?: весь алгоритм не будет быстрее, если я запущу его на ЦП, так как загрузка с ГП занимает слишком много времени (~ 0,00805 с).