Hi, thanks so much for your response. I have revised my code based on your suggestion and I found it seems not to be the problem of the location of the indexing operation, but the operation of torch.nonzero takes up too much time. As convolution is highly sped up by GPU, operations like torch.nonzero seem to be less efficient when the code is running on GPU and this single operation takes almost half of the time of the whole inference process. I wonder whether this observation is normal and is there any way to reduce the time of these operations. Thanks.