feat: improve how device switch is handled between the metric device and the input tensors device

ignite/metrics/ssim.py

@@ -1,3 +1,4 @@
+import warnings
 from typing import Callable, Sequence, Union
 
 import torch
@@ -158,17 +159,31 @@ def update(self, output: Sequence[torch.Tensor]) -> None:
             )
 
         channel = y_pred.size(1)
-        if len(self._kernel.shape) < 4:
-            self._kernel = self._kernel.expand(channel, 1, -1, -1).to(device=y_pred.device)
+
+        self._kernel = self._kernel.expand(channel, 1, -1, -1)
+        kernel = self._kernel
+
+        if y_pred.device != self._kernel.device:
+            warnings.warn(
+                "The metric's device should be the same than your update tensors. See SSIM() device argument.",
+                RuntimeWarning,
+            )
+
+            if self._kernel.device == torch.device("cpu"):
+                kernel = kernel.to(device=y_pred.device)
+
+            if y_pred.device == torch.device("cpu"):
+                y_pred = y_pred.to(device=self._kernel.device)
+                y = y.to(device=self._kernel.device)
 
         y_pred = F.pad(y_pred, [self.pad_w, self.pad_w, self.pad_h, self.pad_h], mode="reflect")
         y = F.pad(y, [self.pad_w, self.pad_w, self.pad_h, self.pad_h], mode="reflect")
 
-        if y_pred.dtype != self._kernel.dtype:
-            self._kernel = self._kernel.to(dtype=y_pred.dtype)
+        if y_pred.dtype != kernel.dtype:
+            kernel = kernel.to(dtype=y_pred.dtype)
 
         input_list = [y_pred, y, y_pred * y_pred, y * y, y_pred * y]
-        outputs = F.conv2d(torch.cat(input_list), self._kernel, groups=channel)
+        outputs = F.conv2d(torch.cat(input_list), kernel, groups=channel)
         batch_size = y_pred.size(0)
         output_list = [outputs[x * batch_size : (x + 1) * batch_size] for x in range(len(input_list))]
 

tests/ignite/metrics/test_ssim.py

@@ -1,3 +1,5 @@
+import warnings
+
 import numpy as np
 import pytest
 import torch
@@ -102,6 +104,38 @@ def test_ssim(
     assert np.allclose(ignite_ssim, skimg_ssim, atol=precision)
 
 
+@pytest.mark.parametrize(
+    "metric_device, y_pred_device",
+    [
+        [torch.device("cpu"), torch.device("cpu")],
+        [torch.device("cpu"), torch.device("cuda")],
+        [torch.device("cuda"), torch.device("cpu")],
+        [torch.device("cuda"), torch.device("cuda")],
+    ],
+)
+def test_ssim_device(available_device, metric_device, y_pred_device):
+    if available_device == "cpu":
+        pytest.skip("This test requires a cuda device.")
+
+    data_range = 1.0
+    sigma = 1.5
+    shape = (12, 5, 256, 256)
+
+    ssim = SSIM(data_range=data_range, sigma=sigma, device=metric_device)
+
+    y_pred = torch.rand(shape, device=y_pred_device)
+    y = y_pred * 0.8
+
+    with pytest.warns() as record_warning:
+        warnings.warn("Avoid pytest DID NOT WARN failure.")
+        ssim.update((y_pred, y))
+
+    if metric_device == y_pred_device:
+        assert len(record_warning) == 1
+    else:
+        assert len(record_warning) == 2
+
+
 def test_ssim_variable_batchsize(available_device):
     # Checks https://github.com/pytorch/ignite/issues/2532
     sigma = 1.5