Repository: Project-MONAI/MONAI

Issue:
enhance the sliding window devices
when `inputs` is on a cuda device:
```py
inputs = data_type(torch.ones((3, 16, 15, 7))).to(device="cuda")
```
and the sliding window is defined with (`device="cpu"`):
```py
output = sliding_window_inference(inputs, ..., sw_device=None, device="cpu")
```

the `output` should be on `device="cpu"`. But it's currently implemented as:
https://github.com/Project-MONAI/MONAI/blob/e37de6913d1d0734c765c9bdcdadb645d313bef8/monai/inferers/utils.py#L282

---
to reproduce the issue:

```py
>>> import torch
>>> from monai.inferers import sliding_window_inference
>>> from monai.data import MetaTensor
>>> a = MetaTensor(torch.zeros(1,1,5,5)).to("cuda")
>>> sliding_window_inference(a, (3,3), 1, lambda x: x, sw_device="cpu", device="cpu").device
device(type='cuda', index=0)
```

the final device should be 'cpu'


Discussion:

Resolving patch:
diff --git a/monai/inferers/utils.py b/monai/inferers/utils.py
--- a/monai/inferers/utils.py
+++ b/monai/inferers/utils.py
@@ -279,7 +279,7 @@ def sliding_window_inference(
     final_output = final_output[0] if is_tensor_output else final_output
 
     if isinstance(inputs, MetaTensor):
-        final_output = convert_to_dst_type(final_output, inputs)[0]  # type: ignore
+        final_output = convert_to_dst_type(final_output, inputs, device=device)[0]  # type: ignore
     return final_output
 
 
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
@@ -243,7 +243,7 @@ def convert_to_cupy(data, dtype: Optional[np.dtype] = None, wrap_sequence: bool
 def convert_data_type(
     data: Any,
     output_type: Optional[Type[NdarrayTensor]] = None,
-    device: Optional[torch.device] = None,
+    device: Union[None, str, torch.device] = None,
     dtype: Union[DtypeLike, torch.dtype] = None,
     wrap_sequence: bool = False,
 ) -> Tuple[NdarrayTensor, type, Optional[torch.device]]:
@@ -307,7 +307,11 @@ def convert_data_type(
 
 
 def convert_to_dst_type(
-    src: Any, dst: NdarrayTensor, dtype: Union[DtypeLike, torch.dtype, None] = None, wrap_sequence: bool = False
+    src: Any,
+    dst: NdarrayTensor,
+    dtype: Union[DtypeLike, torch.dtype, None] = None,
+    wrap_sequence: bool = False,
+    device: Union[None, str, torch.device] = None,
 ) -> Tuple[NdarrayTensor, type, Optional[torch.device]]:
     """
     Convert source data to the same data type and device as the destination data.
@@ -321,12 +325,13 @@ def convert_to_dst_type(
         dtype: an optional argument if the target `dtype` is different from the original `dst`'s data 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.
 
     See Also:
         :func:`convert_data_type`
     """
 
-    device = dst.device if isinstance(dst, torch.Tensor) else None
+    device = dst.device if device is None and isinstance(dst, torch.Tensor) else device
     if dtype is None:
         dtype = dst.dtype
 


Test patch:
diff --git a/tests/test_sliding_window_inference.py b/tests/test_sliding_window_inference.py
--- a/tests/test_sliding_window_inference.py
+++ b/tests/test_sliding_window_inference.py
@@ -9,6 +9,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import itertools
 import unittest
 
 import numpy as np
@@ -85,20 +86,20 @@ def compute(data):
         expected_val = np.ones((1, 3, 16, 15, 7), dtype=np.float32) + 1
         np.testing.assert_allclose(result.cpu().numpy(), expected_val)
 
-    @parameterized.expand([[x] for x in TEST_TORCH_AND_META_TENSORS])
+    @parameterized.expand(list(itertools.product(TEST_TORCH_AND_META_TENSORS, ("cpu", "cuda"), ("cpu", "cuda", None))))
     @skip_if_no_cuda
-    def test_sw_device(self, data_type):
-        inputs = data_type(torch.ones((3, 16, 15, 7))).to(device="cpu")
+    def test_sw_device(self, data_type, device, sw_device):
+        inputs = data_type(torch.ones((3, 16, 15, 7))).to(device=device)
         inputs = list_data_collate([inputs])  # make a proper batch
         roi_shape = (4, 10, 7)
         sw_batch_size = 10
 
         def compute(data):
-            self.assertEqual(data.device.type, "cuda")
-            return data + torch.tensor(1, device="cuda")
+            self.assertEqual(data.device.type, sw_device or device)
+            return data + torch.tensor(1, device=sw_device or device)
 
-        result = sliding_window_inference(inputs, roi_shape, sw_batch_size, compute, sw_device="cuda")
-        np.testing.assert_string_equal(inputs.device.type, result.device.type)
+        result = sliding_window_inference(inputs, roi_shape, sw_batch_size, compute, sw_device=sw_device, device="cpu")
+        np.testing.assert_string_equal("cpu", result.device.type)
         expected_val = np.ones((1, 3, 16, 15, 7), dtype=np.float32) + 1
         np.testing.assert_allclose(result.cpu().numpy(), expected_val)