def test_isolate_rng_cuda(get_cuda_rng, set_cuda_rng):
    """Test that `include_cuda` controls whether isolate_rng also manages torch.cuda's rng."""
    with isolate_rng(include_cuda=False):
        get_cuda_rng.assert_not_called()
    set_cuda_rng.assert_not_called()

    with isolate_rng(include_cuda=True):
        assert get_cuda_rng.call_count == int(torch.cuda.is_available())
    set_cuda_rng.assert_called_once()
Test that `include_cuda` controls whether isolate_rng also manages torch.cuda's rng.
tests/tests_pytorch/utilities/test_seed.py