Hi,
Did you tried calling torch.cuda.synchronize() juste before this line and check if it is not the sync that takes all the time? Remember that the cuda api is asynchronous and only wait for the computation to be done when you actually get them back to the cpu.