whenever you add benchmarks with CUDA, you need to make sure you call torch.cuda.synchronize() right before collecting the time.
For example:
x = x.cuda()
print('1')
y = x ** 2
torch.cuda.synchronize()
print('2')

In your case, add torch.cuda.synchronize() right before print(‘2’) and print(‘3’) and print(‘4’)