Repository: Project-MONAI/MONAI

Issue:
Intensity overflow in `ShiftIntensity`
**Describe the bug**
Now in `ShiftIntensity`, the default is to use the input dtype as the output dtype. As we discussed in https://github.com/Project-MONAI/MONAI/pull/5535, it will cause an issue when the intensity goes beyond the original range. 
https://github.com/Project-MONAI/MONAI/blob/4cf6cf846468bc325d1695e9a13137f88d31daae/monai/transforms/intensity/array.py#L238

I followed @wyli's suggestion and tried to add `convert_data_dtype` referring to https://github.com/pytorch/vision/blob/ad128b753c7e8cc0c600dfddac22ff48fc73c9d9/torchvision/transforms/functional_tensor.py#L64, but I find that what they did may be different with what we want.

```
>>> a = torch.tensor([175, 175], dtype=torch.uint8)
>>> b = convert_image_dtype(a, dtype=torch.float32)
>>> print(b, b.dtype)
tensor([0.6863, 0.6863]) torch.float32
```

I try to add some warning in this [draft pr](https://github.com/Project-MONAI/MONAI/pull/5620). But that still doesn't solve the problem. Because even if the user sees the warning, the user cannot modify it because the `ShiftIntensity` does not have the `dtype` parameter passed in. And if we add `convert_data_dtype` for each transform, it still doesn't look right, because we change the intensity inside the transform, which may not be what the user wants. 
Like what @ericspod and @myron said in https://github.com/Project-MONAI/MONAI/pull/5535, adding a clip transform is indeed a good solution, but it does not solve the current bug, because the data dtype has been changed inside the transform, that is, the intensity has been changed.
At the moment, I have not thought of a better way than adding a WARNING. 
Do you have any suggestions? cc @wyli @ericspod @myron @Nic-Ma 
test_query_memory doesn't work for certain hardware config
```
======================================================================
ERROR: test_output_str (tests.test_query_memory.TestQueryMemory)
----------------------------------------------------------------------
Traceback (most recent call last):
  File “/tmp/MONAI/tests/test_query_memory.py”, line 19, in test_output_str
    self.assertTrue(isinstance(query_memory(2), str))
  File “/tmp/MONAI/tests/utils.py”, line 728, in query_memory
    free_memory = np.asarray(free_memory, dtype=float).T
ValueError: could not convert string to float: ‘[N/A]’
----------------------------------------------------------------------
```


Discussion:

Resolving patch:
diff --git a/monai/utils/type_conversion.py b/monai/utils/type_conversion.py
--- a/monai/utils/type_conversion.py
+++ b/monai/utils/type_conversion.py
@@ -105,6 +105,7 @@ def convert_to_tensor(
     device: Union[None, str, torch.device] = None,
     wrap_sequence: bool = False,
     track_meta: bool = False,
+    safe: bool = False,
 ):
     """
     Utility to convert the input data to a PyTorch Tensor, if `track_meta` is True, the output will be a `MetaTensor`,
@@ -121,6 +122,9 @@ def convert_to_tensor(
             E.g., `[1, 2]` -> `[tensor(1), tensor(2)]`. If `True`, then `[1, 2]` -> `tensor([1, 2])`.
         track_meta: whether to track the meta information, if `True`, will convert to `MetaTensor`.
             default to `False`.
+        safe: if `True`, then do safe dtype convert when intensity overflow. default to `False`.
+            E.g., `[256, -12]` -> `[tensor(0), tensor(244)]`.
+            If `True`, then `[256, -12]` -> `[tensor(255), tensor(0)]`.
 
     """
 
@@ -138,6 +142,8 @@ def _convert_tensor(tensor, **kwargs):
             return tensor.as_tensor()
         return tensor
 
+    if safe:
+        data = safe_dtype_range(data, dtype)
     dtype = get_equivalent_dtype(dtype, torch.Tensor)
     if isinstance(data, torch.Tensor):
         return _convert_tensor(data).to(dtype=dtype, device=device, memory_format=torch.contiguous_format)
@@ -164,7 +170,7 @@ def _convert_tensor(tensor, **kwargs):
     return data
 
 
-def convert_to_numpy(data, dtype: DtypeLike = None, wrap_sequence: bool = False):
+def convert_to_numpy(data, dtype: DtypeLike = None, wrap_sequence: bool = False, safe: bool = False):
     """
     Utility to convert the input data to a numpy array. If passing a dictionary, list or tuple,
     recursively check every item and convert it to numpy array.
@@ -176,7 +182,11 @@ def convert_to_numpy(data, dtype: DtypeLike = None, wrap_sequence: bool = False)
         dtype: target data type when converting to numpy array.
         wrap_sequence: if `False`, then lists will recursively call this function.
             E.g., `[1, 2]` -> `[array(1), array(2)]`. If `True`, then `[1, 2]` -> `array([1, 2])`.
+        safe: if `True`, then do safe dtype convert when intensity overflow. default to `False`.
+            E.g., `[256, -12]` -> `[array(0), array(244)]`. If `True`, then `[256, -12]` -> `[array(255), array(0)]`.
     """
+    if safe:
+        data = safe_dtype_range(data, dtype)
     if isinstance(data, torch.Tensor):
         data = np.asarray(data.detach().to(device="cpu").numpy(), dtype=get_equivalent_dtype(dtype, np.ndarray))
     elif has_cp and isinstance(data, cp_ndarray):
@@ -205,7 +215,7 @@ def convert_to_numpy(data, dtype: DtypeLike = None, wrap_sequence: bool = False)
     return data
 
 
-def convert_to_cupy(data, dtype: Optional[np.dtype] = None, wrap_sequence: bool = False):
+def convert_to_cupy(data, dtype: Optional[np.dtype] = None, wrap_sequence: bool = False, safe: bool = False):
     """
     Utility to convert the input data to a cupy array. If passing a dictionary, list or tuple,
     recursively check every item and convert it to cupy array.
@@ -218,8 +228,11 @@ def convert_to_cupy(data, dtype: Optional[np.dtype] = None, wrap_sequence: bool
             for more details: https://docs.cupy.dev/en/stable/reference/generated/cupy.array.html.
         wrap_sequence: if `False`, then lists will recursively call this function.
             E.g., `[1, 2]` -> `[array(1), array(2)]`. If `True`, then `[1, 2]` -> `array([1, 2])`.
+        safe: if `True`, then do safe dtype convert when intensity overflow. default to `False`.
+            E.g., `[256, -12]` -> `[array(0), array(244)]`. If `True`, then `[256, -12]` -> `[array(255), array(0)]`.
     """
-
+    if safe:
+        data = safe_dtype_range(data, dtype)
     # direct calls
     if isinstance(data, (cp_ndarray, np.ndarray, torch.Tensor, float, int, bool)):
         data = cp.asarray(data, dtype)
@@ -246,6 +259,7 @@ def convert_data_type(
     device: Union[None, str, torch.device] = None,
     dtype: Union[DtypeLike, torch.dtype] = None,
     wrap_sequence: bool = False,
+    safe: bool = False,
 ) -> Tuple[NdarrayTensor, type, Optional[torch.device]]:
     """
     Convert to `MetaTensor`, `torch.Tensor` or `np.ndarray` from `MetaTensor`, `torch.Tensor`,
@@ -260,6 +274,8 @@ def convert_data_type(
             If left blank, it remains unchanged.
         wrap_sequence: if `False`, then lists will recursively call this function.
             E.g., `[1, 2]` -> `[array(1), array(2)]`. If `True`, then `[1, 2]` -> `array([1, 2])`.
+        safe: if `True`, then do safe dtype convert when intensity overflow. default to `False`.
+            E.g., `[256, -12]` -> `[array(0), array(244)]`. If `True`, then `[256, -12]` -> `[array(255), array(0)]`.
 
     Returns:
         modified data, orig_type, orig_device
@@ -288,20 +304,20 @@ def convert_data_type(
     orig_device = data.device if isinstance(data, torch.Tensor) else None
 
     output_type = output_type or orig_type
-
     dtype_ = get_equivalent_dtype(dtype, output_type)
 
     data_: NdarrayTensor
-
     if issubclass(output_type, torch.Tensor):
         track_meta = issubclass(output_type, monai.data.MetaTensor)
-        data_ = convert_to_tensor(data, dtype=dtype_, device=device, wrap_sequence=wrap_sequence, track_meta=track_meta)
+        data_ = convert_to_tensor(
+            data, dtype=dtype_, device=device, wrap_sequence=wrap_sequence, track_meta=track_meta, safe=safe
+        )
         return data_, orig_type, orig_device
     if issubclass(output_type, np.ndarray):
-        data_ = convert_to_numpy(data, dtype=dtype_, wrap_sequence=wrap_sequence)
+        data_ = convert_to_numpy(data, dtype=dtype_, wrap_sequence=wrap_sequence, safe=safe)
         return data_, orig_type, orig_device
     elif has_cp and issubclass(output_type, cp.ndarray):
-        data_ = convert_to_cupy(data, dtype=dtype_, wrap_sequence=wrap_sequence)
+        data_ = convert_to_cupy(data, dtype=dtype_, wrap_sequence=wrap_sequence, safe=safe)
         return data_, orig_type, orig_device
     raise ValueError(f"Unsupported output type: {output_type}")
 
@@ -312,6 +328,7 @@ def convert_to_dst_type(
     dtype: Union[DtypeLike, torch.dtype, None] = None,
     wrap_sequence: bool = False,
     device: Union[None, str, torch.device] = None,
+    safe: bool = False,
 ) -> Tuple[NdarrayTensor, type, Optional[torch.device]]:
     """
     Convert source data to the same data type and device as the destination data.
@@ -326,6 +343,8 @@ def convert_to_dst_type(
         wrap_sequence: if `False`, then lists will recursively call this function. E.g., `[1, 2]` -> `[array(1), array(2)]`.
             If `True`, then `[1, 2]` -> `array([1, 2])`.
         device: target device to put the converted Tensor data. If unspecified, `dst.device` will be used if possible.
+        safe: if `True`, then do safe dtype convert when intensity overflow. default to `False`.
+            E.g., `[256, -12]` -> `[array(0), array(244)]`. If `True`, then `[256, -12]` -> `[array(255), array(0)]`.
 
     See Also:
         :func:`convert_data_type`
@@ -349,7 +368,7 @@ def convert_to_dst_type(
         output_type = type(dst)
     output: NdarrayTensor
     output, _type, _device = convert_data_type(
-        data=src, output_type=output_type, device=device, dtype=dtype, wrap_sequence=wrap_sequence
+        data=src, output_type=output_type, device=device, dtype=dtype, wrap_sequence=wrap_sequence, safe=safe
     )
     if copy_meta and isinstance(output, monai.data.MetaTensor):
         output.copy_meta_from(dst)
@@ -366,3 +385,74 @@ def convert_to_list(data: Union[Sequence, torch.Tensor, np.ndarray]) -> list:
 
     """
     return data.tolist() if isinstance(data, (torch.Tensor, np.ndarray)) else list(data)
+
+
+def get_dtype_bound_value(dtype: Union[DtypeLike, torch.dtype]):
+    """
+    Get dtype bound value
+    Args:
+        dtype: dtype to get bound value
+    Returns:
+        (bound_min_value, bound_max_value)
+    """
+    if dtype in UNSUPPORTED_TYPES:
+        is_floating_point = False
+    else:
+        is_floating_point = get_equivalent_dtype(dtype, torch.Tensor).is_floating_point
+    dtype = get_equivalent_dtype(dtype, np.array)
+    if is_floating_point:
+        return (np.finfo(dtype).min, np.finfo(dtype).max)  # type: ignore
+    else:
+        return (np.iinfo(dtype).min, np.iinfo(dtype).max)
+
+
+def safe_dtype_range(data: Any, dtype: Union[DtypeLike, torch.dtype] = None):
+    """
+    Utility to safely convert the input data to target dtype.
+
+    Args:
+        data: input data can be PyTorch Tensor, numpy array, list, dictionary, int, float, bool, str, etc.
+            will convert to target dtype and keep the original type.
+            for dictionary, list or tuple, convert every item.
+        dtype: target data type to convert.
+    """
+
+    def _safe_dtype_range(data, dtype):
+        output_dtype = dtype if dtype is not None else data.dtype
+        dtype_bound_value = get_dtype_bound_value(output_dtype)
+        if data.ndim == 0:
+            data_bound = (data, data)
+        else:
+            data_bound = (min(data), max(data))
+        if (data_bound[1] > dtype_bound_value[1]) or (data_bound[0] < dtype_bound_value[0]):
+            if isinstance(data, torch.Tensor):
+                return torch.clamp(data, dtype_bound_value[0], dtype_bound_value[1])
+            elif isinstance(data, np.ndarray):
+                return np.clip(data, dtype_bound_value[0], dtype_bound_value[1])
+            elif has_cp and isinstance(data, cp_ndarray):
+                return cp.clip(data, dtype_bound_value[0], dtype_bound_value[1])
+        else:
+            return data
+
+    if has_cp and isinstance(data, cp_ndarray):
+        return cp.asarray(_safe_dtype_range(data, dtype))
+    elif isinstance(data, np.ndarray):
+        return np.asarray(_safe_dtype_range(data, dtype))
+    elif isinstance(data, torch.Tensor):
+        return _safe_dtype_range(data, dtype)
+    elif isinstance(data, (float, int, bool)) and dtype is None:
+        return data
+    elif isinstance(data, (float, int, bool)) and dtype is not None:
+        output_dtype = dtype
+        dtype_bound_value = get_dtype_bound_value(output_dtype)
+        data = dtype_bound_value[1] if data > dtype_bound_value[1] else data
+        data = dtype_bound_value[0] if data < dtype_bound_value[0] else data
+        return data
+
+    elif isinstance(data, list):
+        return [safe_dtype_range(i, dtype=dtype) for i in data]
+    elif isinstance(data, tuple):
+        return tuple(safe_dtype_range(i, dtype=dtype) for i in data)
+    elif isinstance(data, dict):
+        return {k: safe_dtype_range(v, dtype=dtype) for k, v in data.items()}
+    return data


Test patch:
diff --git a/tests/test_convert_data_type.py b/tests/test_convert_data_type.py
--- a/tests/test_convert_data_type.py
+++ b/tests/test_convert_data_type.py
@@ -17,27 +17,45 @@
 from parameterized import parameterized
 
 from monai.data import MetaTensor
-from monai.utils.type_conversion import convert_data_type, convert_to_dst_type
-from tests.utils import TEST_NDARRAYS_ALL
+from monai.utils.type_conversion import convert_data_type, convert_to_dst_type, get_equivalent_dtype
+from tests.utils import TEST_NDARRAYS_ALL, assert_allclose
 
 TESTS: List[Tuple] = []
 for in_type in TEST_NDARRAYS_ALL + (int, float):
     for out_type in TEST_NDARRAYS_ALL:
-        TESTS.append((in_type(np.array(1.0)), out_type(np.array(1.0))))  # type: ignore
+        TESTS.append((in_type(np.array(1.0)), out_type(np.array(1.0)), None, False))  # type: ignore
+        if in_type is not float:
+            TESTS.append((in_type(np.array(256)), out_type(np.array(255)), np.uint8, True))  # type: ignore
 
 TESTS_LIST: List[Tuple] = []
 for in_type in TEST_NDARRAYS_ALL + (int, float):
     for out_type in TEST_NDARRAYS_ALL:
         TESTS_LIST.append(
-            ([in_type(np.array(1.0)), in_type(np.array(1.0))], out_type(np.array([1.0, 1.0])), True)  # type: ignore
+            ([in_type(np.array(1.0)), in_type(np.array(1.0))], out_type(np.array([1.0, 1.0])), True, None, False)  # type: ignore
         )
         TESTS_LIST.append(
             (
                 [in_type(np.array(1.0)), in_type(np.array(1.0))],  # type: ignore
                 [out_type(np.array(1.0)), out_type(np.array(1.0))],
                 False,
+                None,
+                False,
             )
         )
+        if in_type is not float:
+            TESTS_LIST.append(
+                ([in_type(np.array(257)), in_type(np.array(1))], out_type(np.array([255, 1])), True, np.uint8, True)  # type: ignore
+            )
+            TESTS_LIST.append(
+                (
+                    [in_type(np.array(257)), in_type(np.array(-12))],  # type: ignore
+                    [out_type(np.array(255)), out_type(np.array(0))],
+                    False,
+                    np.uint8,
+                    True,
+                )
+            )
+
 
 UNSUPPORTED_TYPES = {np.dtype("uint16"): torch.int32, np.dtype("uint32"): torch.int64, np.dtype("uint64"): torch.int64}
 
@@ -48,17 +66,20 @@ class TestTensor(torch.Tensor):
 
 class TestConvertDataType(unittest.TestCase):
     @parameterized.expand(TESTS)
-    def test_convert_data_type(self, in_image, im_out):
-        converted_im, orig_type, orig_device = convert_data_type(in_image, type(im_out))
+    def test_convert_data_type(self, in_image, im_out, out_dtype, safe):
+        converted_im, orig_type, orig_device = convert_data_type(in_image, type(im_out), dtype=out_dtype, safe=safe)
         # check input is unchanged
         self.assertEqual(type(in_image), orig_type)
         if isinstance(in_image, torch.Tensor):
             self.assertEqual(in_image.device, orig_device)
         # check output is desired type
         self.assertEqual(type(converted_im), type(im_out))
+        # check data has been clipped
+        assert_allclose(converted_im, im_out)
         # check dtype is unchanged
-        if isinstance(in_type, (np.ndarray, torch.Tensor)):
-            self.assertEqual(converted_im.dtype, im_out.dtype)
+        if out_dtype is None:
+            if isinstance(in_image, (np.ndarray, torch.Tensor)):
+                self.assertEqual(converted_im.dtype, im_out.dtype)
 
     def test_neg_stride(self):
         _ = convert_data_type(np.array((1, 2))[::-1], torch.Tensor)
@@ -71,26 +92,32 @@ def test_unsupported_np_types(self, np_type, pt_type):
         self.assertEqual(converted_im.dtype, pt_type)
 
     @parameterized.expand(TESTS_LIST)
-    def test_convert_list(self, in_image, im_out, wrap):
+    def test_convert_list(self, in_image, im_out, wrap, out_dtype, safe):
         output_type = type(im_out) if wrap else type(im_out[0])
-        converted_im, *_ = convert_data_type(in_image, output_type, wrap_sequence=wrap)
+        converted_im, *_ = convert_data_type(in_image, output_type, wrap_sequence=wrap, dtype=out_dtype, safe=safe)
         # check output is desired type
         if not wrap:
             converted_im = converted_im[0]
             im_out = im_out[0]
         self.assertEqual(type(converted_im), type(im_out))
+        assert_allclose(converted_im, im_out)
         # check dtype is unchanged
-        if isinstance(in_type, (np.ndarray, torch.Tensor)):
-            self.assertEqual(converted_im.dtype, im_out.dtype)
+        if isinstance(in_image[0], (np.ndarray, torch.Tensor)):
+            if out_dtype is None:
+                self.assertEqual(converted_im.dtype, im_out.dtype)
+            else:
+                _out_dtype = get_equivalent_dtype(out_dtype, output_type)
+                self.assertEqual(converted_im.dtype, _out_dtype)
 
 
 class TestConvertDataSame(unittest.TestCase):
     # add test for subclass of Tensor
-    @parameterized.expand(TESTS + [(np.array(1.0), TestTensor(np.array(1.0)))])
-    def test_convert_data_type(self, in_image, im_out):
-        converted_im, orig_type, orig_device = convert_to_dst_type(in_image, im_out)
+    @parameterized.expand(TESTS + [(np.array(256), TestTensor(np.array([255])), torch.uint8, True)])
+    def test_convert_data_type(self, in_image, im_out, out_dtype, safe):
+        converted_im, orig_type, orig_device = convert_to_dst_type(in_image, im_out, dtype=out_dtype, safe=safe)
         # check input is unchanged
         self.assertEqual(type(in_image), orig_type)
+        assert_allclose(converted_im, im_out)
         if isinstance(in_image, torch.Tensor):
             self.assertEqual(in_image.device, orig_device)
 
@@ -103,8 +130,9 @@ def test_convert_data_type(self, in_image, im_out):
             output_type = np.ndarray
         self.assertEqual(type(converted_im), output_type)
         # check dtype is unchanged
-        if isinstance(in_type, (np.ndarray, torch.Tensor, MetaTensor)):
-            self.assertEqual(converted_im.dtype, im_out.dtype)
+        if out_dtype is None:
+            if isinstance(in_image, (np.ndarray, torch.Tensor, MetaTensor)):
+                self.assertEqual(converted_im.dtype, im_out.dtype)
 
 
 if __name__ == "__main__":
diff --git a/tests/test_safe_dtype_range.py b/tests/test_safe_dtype_range.py
new file mode 100644
--- /dev/null
+++ b/tests/test_safe_dtype_range.py
@@ -0,0 +1,96 @@
+# Copyright (c) MONAI Consortium
+# 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.
+
+import unittest
+from typing import List, Tuple
+
+import numpy as np
+import torch
+from parameterized import parameterized
+
+from monai.utils import optional_import
+from monai.utils.type_conversion import get_equivalent_dtype, safe_dtype_range
+from tests.utils import HAS_CUPY, TEST_NDARRAYS_ALL, assert_allclose
+
+cp, _ = optional_import("cupy")
+
+TESTS: List[Tuple] = []
+for in_type in TEST_NDARRAYS_ALL + (int, float):
+    TESTS.append((in_type(np.array(1.0)), in_type(np.array(1.0)), None))  # type: ignore
+    if in_type is not float:
+        TESTS.append((in_type(np.array(256)), in_type(np.array(255)), np.uint8))  # type: ignore
+        TESTS.append((in_type(np.array(-12)), in_type(np.array(0)), np.uint8))  # type: ignore
+
+TESTS_LIST: List[Tuple] = []
+for in_type in TEST_NDARRAYS_ALL + (int, float):
+    TESTS_LIST.append(
+        (
+            [in_type(np.array(1.0)), in_type(np.array(1.0))],  # type: ignore
+            [in_type(np.array(1.0)), in_type(np.array(1.0))],  # type: ignore
+            None,
+        )
+    )
+    if in_type is not float:
+        TESTS_LIST.append(
+            (
+                [in_type(np.array(257)), in_type(np.array(-12))],  # type: ignore
+                [in_type(np.array(255)), in_type(np.array(0))],  # type: ignore
+                np.uint8,
+            )
+        )
+
+TESTS_CUPY = [[np.array(1.0), np.array(1.0), None], [np.array([-12]), np.array([0]), np.uint8]]
+
+
+class TesSafeDtypeRange(unittest.TestCase):
+    @parameterized.expand(TESTS)
+    def test_safe_dtype_range(self, in_image, im_out, out_dtype):
+        result = safe_dtype_range(in_image, out_dtype)
+        # check type is unchanged
+        self.assertEqual(type(in_image), type(result))
+        # check dtype is unchanged
+        if isinstance(in_image, (np.ndarray, torch.Tensor)):
+            self.assertEqual(in_image.dtype, result.dtype)
+        # check output
+        assert_allclose(result, im_out)
+
+    @parameterized.expand(TESTS_LIST)
+    def test_safe_dtype_range_list(self, in_image, im_out, out_dtype):
+        output_type = type(im_out[0])
+        result = safe_dtype_range(in_image, dtype=out_dtype)
+        # check type is unchanged
+        self.assertEqual(type(result), type(im_out))
+        # check output
+        for i, _result in enumerate(result):
+            assert_allclose(_result, im_out[i])
+        # check dtype is unchanged
+        if isinstance(in_image, (np.ndarray, torch.Tensor)):
+            if out_dtype is None:
+                self.assertEqual(result[0].dtype, im_out[0].dtype)
+            else:
+                _out_dtype = get_equivalent_dtype(out_dtype, output_type)
+                self.assertEqual(result[0].dtype, _out_dtype)
+
+    @parameterized.expand(TESTS_CUPY)
+    @unittest.skipUnless(HAS_CUPY, "Requires CuPy")
+    def test_type_cupy(self, in_image, im_out, out_dtype):
+        in_image = cp.asarray(in_image)
+        result = safe_dtype_range(in_image, dtype=out_dtype)
+        # check type is unchanged
+        self.assertEqual(type(in_image), type(result))
+        # check dtype is unchanged
+        self.assertEqual(result.dtype, in_image.dtype)
+        # check output
+        self.assertEqual(result, cp.asarray(im_out))
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/tests/utils.py b/tests/utils.py
--- a/tests/utils.py
+++ b/tests/utils.py
@@ -744,7 +744,7 @@ def query_memory(n=2):
         free_memory = np.asarray(free_memory, dtype=float).T
         free_memory[1] += free_memory[0]  # combine 0/1 column measures
         ids = np.lexsort(free_memory)[:n]
-    except (TypeError, IndexError, OSError):
+    except (TypeError, ValueError, IndexError, OSError):
         ids = range(n) if isinstance(n, int) else []
     return ",".join(f"{int(x)}" for x in ids)