I got more clear code below after I’ve read the doc. And I find that the performance problem comes from ‘nonzero’ function
mask = Y[i].nonzero().squeeze()
x = torch.index_select(x, 0, mask)

Then I found the same issue on google:
torch.nonzero slower than np.nonzero · Issue #14848 · pytorch/pytorch (github.com) 28
GPU version ‘nonzero’ function is still much slower than that of CPU version