diff --git a/monai/utils/type_conversion.py b/monai/utils/type_conversion.py index b0ce187e38..3688b02d26 100644 --- a/monai/utils/type_conversion.py +++ b/monai/utils/type_conversion.py @@ -105,6 +105,8 @@ def convert_to_tensor(data, wrap_sequence: bool = False): # numpy array with 0 dims is also sequence iterable, # `ascontiguousarray` will add 1 dim if img has no dim, so we only apply on data with dims return torch.as_tensor(data if data.ndim == 0 else np.ascontiguousarray(data)) + elif has_cp and isinstance(data, cp_ndarray): + return torch.as_tensor(data) elif isinstance(data, (float, int, bool)): return torch.as_tensor(data) elif isinstance(data, Sequence) and wrap_sequence: diff --git a/tests/test_to_tensor.py b/tests/test_to_tensor.py index 6ac06983f6..3d187a1dba 100644 --- a/tests/test_to_tensor.py +++ b/tests/test_to_tensor.py @@ -12,9 +12,12 @@ import unittest from parameterized import parameterized +from torch import Tensor from monai.transforms import ToTensor -from tests.utils import TEST_NDARRAYS, assert_allclose +from tests.utils import TEST_NDARRAYS, assert_allclose, optional_import + +cp, has_cp = optional_import("cupy") im = [[1, 2], [3, 4]] @@ -33,15 +36,25 @@ class TestToTensor(unittest.TestCase): @parameterized.expand(TESTS) def test_array_input(self, test_data, expected_shape): result = ToTensor()(test_data) + self.assertTrue(isinstance(result, Tensor)) assert_allclose(result, test_data) self.assertTupleEqual(result.shape, expected_shape) @parameterized.expand(TESTS_SINGLE) def test_single_input(self, test_data): result = ToTensor()(test_data) + self.assertTrue(isinstance(result, Tensor)) assert_allclose(result, test_data) self.assertEqual(result.ndim, 0) + @unittest.skipUnless(has_cp, "CuPy is required.") + def test_cupy(self): + test_data = [[1, 2], [3, 4]] + cupy_array = cp.ascontiguousarray(cp.asarray(test_data)) + result = ToTensor()(cupy_array) + self.assertTrue(isinstance(result, Tensor)) + assert_allclose(result, test_data) + if __name__ == "__main__": unittest.main()