Hi,
thank you for your answer!

Also maybe you want to remove all the data initialization from the benchmark? Or do you actually care about that?

I somewhat care about this. But as we do ~86000 loops the initialization shouldn’t have any impact on the endresult.

Also code-wise, the problem is that your two implementation do different things so that might explain the perf difference between the two.
Have you tries running the numpy code with pytorch directly (the content of the for loop). All these advanced indexing ops should work properly afaik

I think they are doing the same thing. I copied the pytorch code and replaced it with the numpy functionality. I also added verification that the endresults are the same for all methods. (test_pytorch_implementation_correctness and test_cupy_implementation_correctness)
Copying the “numpy loop” over makes the results much worse (only tested on cpu):

TorchScript 15s (N=500)/ 77s(N=10000)
pytorch 24s (N=500) / 87s (N=10000)

This fits with my previous experience that using the pytorch functions is a lot faster than the python operations.

I am not sure how the pytest benchmark utilities are working but they are most likely not doing proper cuda synchronization. So you most likely want to add the proper calls there (for pytorch torch.cuda.synchronize(), not sure for cupy).

I will check this, thanks!