diff --git a/src/lerobot/utils/device_utils.py b/src/lerobot/utils/device_utils.py index 3f9b58773..1ac301a47 100644 --- a/src/lerobot/utils/device_utils.py +++ b/src/lerobot/utils/device_utils.py @@ -37,16 +37,25 @@ def auto_select_torch_device() -> torch.device: # TODO(Steven): Remove log. log shouldn't be an argument, this should be handled by the logger level def get_safe_torch_device(try_device: str, log: bool = False) -> torch.device: - """Given a string, return a torch.device with checks on whether the device is available.""" + """Given a string, return a torch.device with checks on whether the device is available. + + Raises: + ValueError: If the requested device family is known but not available on + this machine (``AssertionError`` was previously used and is easy to + mistake for a programmer bug under ``python -O`` where asserts vanish). + """ try_device = str(try_device) if try_device.startswith("cuda"): - assert torch.cuda.is_available() + if not torch.cuda.is_available(): + raise ValueError(f"Requested device {try_device!r} but CUDA is not available.") device = torch.device(try_device) elif try_device == "mps": - assert torch.backends.mps.is_available() + if not torch.backends.mps.is_available(): + raise ValueError("Requested device 'mps' but MPS is not available.") device = torch.device("mps") elif try_device == "xpu": - assert torch.xpu.is_available() + if not torch.xpu.is_available(): + raise ValueError("Requested device 'xpu' but XPU is not available.") device = torch.device("xpu") elif try_device == "cpu": device = torch.device("cpu") diff --git a/tests/utils/test_device_utils.py b/tests/utils/test_device_utils.py new file mode 100644 index 000000000..ed62af397 --- /dev/null +++ b/tests/utils/test_device_utils.py @@ -0,0 +1,36 @@ +#!/usr/bin/env python +# Copyright 2026 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from unittest.mock import patch + +import pytest +import torch + +from lerobot.utils.device_utils import get_safe_torch_device, is_torch_device_available + + +def test_cpu_always_available(): + assert get_safe_torch_device("cpu") == torch.device("cpu") + assert is_torch_device_available("cpu") + + +def test_missing_cuda_raises_valueerror(): + with patch("torch.cuda.is_available", return_value=False), pytest.raises(ValueError, match="CUDA"): + get_safe_torch_device("cuda") + + +def test_missing_mps_raises_valueerror(): + with patch("torch.backends.mps.is_available", return_value=False), pytest.raises(ValueError, match="MPS"): + get_safe_torch_device("mps")