Request torch.cuda RNG states only if CUDA is available (#19234)

Author: awaelchli

src/lightning/fabric/utilities/seed.py

@@ -115,7 +115,7 @@ def _collect_rng_states(include_cuda: bool = True) -> Dict[str, Any]:
         "python": python_get_rng_state(),
     }
     if include_cuda:
-        states["torch.cuda"] = torch.cuda.get_rng_state_all()
+        states["torch.cuda"] = torch.cuda.get_rng_state_all() if torch.cuda.is_available() else []
     return states
 
 

tests/tests_fabric/utilities/test_seed.py

@@ -1,5 +1,6 @@
 import os
 from unittest import mock
+from unittest.mock import Mock
 
 import lightning.fabric.utilities
 import pytest
@@ -81,3 +82,12 @@ def test_backward_compatibility_rng_states_dict():
     assert "torch.cuda" in states
     states.pop("torch.cuda")
     _set_rng_states(states)
+
+
+@mock.patch("lightning.fabric.utilities.seed.torch.cuda.is_available", Mock(return_value=False))
+@mock.patch("lightning.fabric.utilities.seed.torch.cuda.get_rng_state_all")
+def test_collect_rng_states_if_cuda_init_fails(get_rng_state_all_mock):
+    """Test that the `torch.cuda` rng states are only requested if CUDA is available."""
+    get_rng_state_all_mock.side_effect = RuntimeError("The NVIDIA driver on your system is too old")
+    states = _collect_rng_states()
+    assert states["torch.cuda"] == []

tests/tests_pytorch/utilities/test_seed.py

@@ -47,5 +47,5 @@ def test_isolate_rng_cuda(get_cuda_rng, set_cuda_rng):
     set_cuda_rng.assert_not_called()
 
     with isolate_rng(include_cuda=True):
-        get_cuda_rng.assert_called_once()
+        assert get_cuda_rng.call_count == int(torch.cuda.is_available())
     set_cuda_rng.assert_called_once()