Solved by lelouedec in post #2 
              
              
                Oops I just found the solution on the documentation under a part I haven’t seen before : 
https://pytorch.org/docs/stable/cuda.html 
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.