Hi,
Why do you think the transfer is slow?
Keep in mind that the cuda API is asynchronous except when it needs to deal with CPU values. So if you measure without manual synchronization with torch.cuda.synchronize(), then it will appear that the copies are slow only because they wait for the rest of the computations before being able to execute as they deal with CPU values.