at::cuda::CUDAStreamGuard guard(stream) is what you are looking for. Example:

  
      github.com
  
  
    pytorch/pytorch/blob/13e6326c07cd44ad52db32ac3918eeb79f8230bc/aten/src/ATen/test/cuda_stream_test.cpp#L152-L162 42

// Setting a stream changes the current device and the stream on that device
{
  at::cuda::CUDAStreamGuard guard(streams1[1]);
  ASSERT_EQ_CUDA(guard.current_device(), at::Device(at::kCUDA, 1));
  ASSERT_EQ_CUDA(at::cuda::current_device(), 1);
  ASSERT_EQ_CUDA(at::cuda::getCurrentCUDAStream(1), streams1[1]);
}


// Device and stream are now reset
ASSERT_EQ_CUDA(at::cuda::current_device(), 0);
ASSERT_EQ_CUDA(at::cuda::getCurrentCUDAStream(1), streams1[0]);