torch.cuda.synchronize
torch.cuda.synchronize(device=None)
Waits for all kernels in all streams on a CUDA device to complete.
Parameters:
      device (torch.device or int, optional) -- device
      for which to synchronize. It uses the current device, given by
      "current_device()", if "device" is "None" (default).