Oops I just found the solution on the documentation under a part I haven’t seen before :
https://pytorch.org/docs/stable/cuda.html 44
For any body who wants to wait for the GPU code to execute you can use :

torch.cuda.synchronize()

Which will wait for the current kernel to end on the GPU.