Hi,
I guess you run this on gpu?
The point is that the cuda api is asynchronous, and to speed things up as much as possible, pytorch will not block while operations are running on gpu unless you ask for the result on cpu side.
I guess your criterion here is the first place where things are needed on the cpu side and so it is waiting for all operations that were queued on the GPU to finish running.
To get accurate timings while running on GPU, you would need to use torch.cuda.synchronize() to explicitly wait for operations on GPU to finish. Or with the debugger you’re using, use the CUDA_LAUNCH_BLOCKING=1 env variable from CUDA (you can look online for more details), but this can affect the performances quite a bit and I’m not sure this will give proper timings.