Hi,
Keep in mind that cuda api is asynchronous and will only synchronize when it interacts with the cpu. So I guess the backward call is super fast but because it’s only the time to queue the work onto the gpu. But as soon as you want to access this information on the cpu (or send data to the gpu), you have to wait for all these computations to be done. So the long runtime for the .cuda() that you see is most certainly the time to actually compute your backward. You can use torch.cuda.synchronize() to force synchronization in your code. Adding one just after the backward call should reduce the runtime of the .cuda() !