Thanks for sharing the code!
Be a bit careful about the CUDA numbers, since you didn’t synchronize the calls.
CUDA calls will be executed asynchronously, which means the CPU can continue executing the code while the CUDA kernels are busy until a synchronization point is reached.
These points can be added manually using torch.cuda.synchronize() are are automatically reached, e.g. if a result from the CUDA operation is needed.
In your profiling script you might in fact just time the kernel launch times, when the .cuda() call is the first op on the tensor.
To proper time CUDA calls, you could have to synchronize before starting and stopping the timer:
torch.cuda.synchronize()
t0 = time.time()
...
torch.cuda.synchronize()
t1 = time.time()