diff --git a/src/lerobot/utils/random_utils.py b/src/lerobot/utils/random_utils.py index e280fc342..6edcf2155 100644 --- a/src/lerobot/utils/random_utils.py +++ b/src/lerobot/utils/random_utils.py @@ -85,6 +85,8 @@ def serialize_torch_rng_state() -> dict[str, torch.Tensor]: torch_rng_state_dict = {"torch_rng_state": torch.get_rng_state()} if torch.cuda.is_available(): torch_rng_state_dict["torch_cuda_rng_state"] = torch.cuda.get_rng_state() + if torch.backends.mps.is_available(): + torch_rng_state_dict["torch_mps_rng_state"] = torch.mps.get_rng_state() return torch_rng_state_dict @@ -95,6 +97,8 @@ def deserialize_torch_rng_state(rng_state_dict: dict[str, torch.Tensor]) -> None torch.set_rng_state(rng_state_dict["torch_rng_state"]) if torch.cuda.is_available() and "torch_cuda_rng_state" in rng_state_dict: torch.cuda.set_rng_state(rng_state_dict["torch_cuda_rng_state"]) + if torch.backends.mps.is_available() and "torch_mps_rng_state" in rng_state_dict: + torch.mps.set_rng_state(rng_state_dict["torch_mps_rng_state"]) def serialize_rng_state() -> dict[str, torch.Tensor]: diff --git a/tests/utils/test_random_utils.py b/tests/utils/test_random_utils.py index e3a5d420f..c79e3d506 100644 --- a/tests/utils/test_random_utils.py +++ b/tests/utils/test_random_utils.py @@ -73,6 +73,17 @@ def test_serialize_deserialize_torch_rng(fixed_seed): assert val2 == val3 +@pytest.mark.skipif(not torch.backends.mps.is_available(), reason="MPS not available") +def test_serialize_deserialize_torch_rng_mps(fixed_seed): + _ = torch.rand(1, device="mps").item() + st = serialize_torch_rng_state() + assert "torch_mps_rng_state" in st + val2 = torch.rand(1, device="mps").item() + deserialize_torch_rng_state(st) + val3 = torch.rand(1, device="mps").item() + assert val2 == val3 + + def test_serialize_deserialize_rng(fixed_seed): # Generate one from each library _ = random.random()