I have a big tensor X whose shape is [?, 18900, 85] and another big tensor Y whose shape is [18900,].
Data of X is float and data of Y is boolean.
The last broadcast operation below takes about 25ms if X and Y both in GPU:
X = YoloModel(image)
# ...
# non max suppresion preprocess
Y = X[..., 4] > conf_thres
for i, x in enumerate(X):
    # x shape is [18900, 85]
    x = x[Y[i]]  # 25ms!!
# X shape now is [47, 85] for example
# ...

Althought it takes less than 1ms when X and Y both reside in CPU, calling X.cpu() would take roughly 25ms.  So I want to know how to fix this problem.
Thanks for your help!