Hi,
I added 3 operations to my code (max/min,  topk,  le/gt).  Besides the topk the other 2 are also used in other locations in my code, though once I added the topk function training became ~x4 slower.
If topk does not move data implicitly to cpu, then how can I speed it up ?   Is there any way to check the implementation of topk in pytorch to see how it distributes the jobs between threads ?
Thanks,
Moshe