5919 fix generalized dice issue (#5929)

Signed-off-by: Yiheng Wang <vennw@nvidia.com>

Fixes #5919 .

### Description

This PR is used to fix the device issue of function
`compute_generalized_dice`, and cuda tensor input will not raise errors.

### Types of changes
<!--- Put an `x` in all the boxes that apply, and remove the not
applicable items -->
- [x] Non-breaking change (fix or new feature that would not break
existing functionality).
- [ ] Breaking change (fix or new feature that would cause existing
functionality to change).
- [ ] New tests added to cover the changes.
- [ ] Integration tests passed locally by running `./runtests.sh -f -u
--net --coverage`.
- [ ] Quick tests passed locally by running `./runtests.sh --quick
--unittests --disttests`.
- [ ] In-line docstrings updated.
- [ ] Documentation updated, tested `make html` command in the `docs/`
folder.

---------

Signed-off-by: Yiheng Wang <vennw@nvidia.com>

Author: yiheng-wang-nv

monai/metrics/generalized_dice.py

@@ -176,7 +176,9 @@ def compute_generalized_dice(
     y_pred_o = y_pred_o.sum(dim=-1)
     denom_zeros = denom == 0
     generalized_dice_score[denom_zeros] = torch.where(
-        (y_pred_o == 0)[denom_zeros], torch.tensor(1.0), torch.tensor(0.0)
+        (y_pred_o == 0)[denom_zeros],
+        torch.tensor(1.0, device=generalized_dice_score.device),
+        torch.tensor(0.0, device=generalized_dice_score.device),
     )
 
     return generalized_dice_score

tests/test_compute_confusion_matrix.py

@@ -14,6 +14,7 @@
 import unittest
 from typing import Any
 
+import numpy as np
 import torch
 from parameterized import parameterized
 
@@ -220,6 +221,7 @@ def test_value(self, input_data, expected_value):
         input_data["include_background"] = False
         result = get_confusion_matrix(**input_data)
         assert_allclose(result, expected_value[:, 1:, :], atol=1e-4, rtol=1e-4)
+        np.testing.assert_equal(result.device, input_data["y_pred"].device)
 
     @parameterized.expand(TEST_CASES_COMPUTE_SAMPLE)
     def test_compute_sample(self, input_data, expected_value):

tests/test_compute_f_beta.py

@@ -13,20 +13,24 @@
 
 import unittest
 
-import numpy
+import numpy as np
 import torch
 
 from monai.metrics import FBetaScore
 from tests.utils import assert_allclose
 
+_device = "cuda:0" if torch.cuda.is_available() else "cpu"
+
 
 class TestFBetaScore(unittest.TestCase):
-    def test_expecting_success(self):
+    def test_expecting_success_and_device(self):
         metric = FBetaScore()
-        metric(
-            y_pred=torch.Tensor([[1, 1, 1], [1, 1, 1], [1, 1, 1]]), y=torch.Tensor([[1, 0, 1], [0, 1, 0], [1, 0, 1]])
-        )
-        assert_allclose(metric.aggregate()[0], torch.Tensor([0.714286]), atol=1e-6, rtol=1e-6)
+        y_pred = torch.tensor([[1, 1, 1], [1, 1, 1], [1, 1, 1]], device=_device)
+        y = torch.tensor([[1, 0, 1], [0, 1, 0], [1, 0, 1]], device=_device)
+        metric(y_pred=y_pred, y=y)
+        result = metric.aggregate()[0]
+        assert_allclose(result, torch.Tensor([0.714286]), atol=1e-6, rtol=1e-6)
+        np.testing.assert_equal(result.device, y_pred.device)
 
     def test_expecting_success2(self):
         metric = FBetaScore(beta=0.5)
@@ -58,7 +62,7 @@ def test_with_nan_values(self):
         metric = FBetaScore(get_not_nans=True)
         metric(
             y_pred=torch.Tensor([[1, 1, 1], [1, 1, 1], [1, 1, 1]]),
-            y=torch.Tensor([[1, 0, 1], [numpy.NaN, numpy.NaN, numpy.NaN], [1, 0, 1]]),
+            y=torch.Tensor([[1, 0, 1], [np.NaN, np.NaN, np.NaN], [1, 0, 1]]),
         )
         assert_allclose(metric.aggregate()[0][0], torch.Tensor([0.727273]), atol=1e-6, rtol=1e-6)
 

tests/test_compute_generalized_dice.py

@@ -19,11 +19,13 @@
 
 from monai.metrics import GeneralizedDiceScore, compute_generalized_dice
 
+_device = "cuda:0" if torch.cuda.is_available() else "cpu"
+
 # keep background
 TEST_CASE_1 = [  # y (1, 1, 2, 2), y_pred (1, 1, 2, 2), expected out (1)
     {
-        "y_pred": torch.tensor([[[[1.0, 0.0], [0.0, 1.0]]]]),
-        "y": torch.tensor([[[[1.0, 0.0], [1.0, 1.0]]]]),
+        "y_pred": torch.tensor([[[[1.0, 0.0], [0.0, 1.0]]]], device=_device),
+        "y": torch.tensor([[[[1.0, 0.0], [1.0, 1.0]]]], device=_device),
         "include_background": True,
     },
     [0.8],
@@ -116,7 +118,12 @@
 TEST_CASE_9 = [{"y": torch.zeros((2, 2, 3, 3)), "y_pred": torch.zeros((2, 2, 3, 3))}, [1.0000, 1.0000]]
 
 
-class TestComputeMeanDice(unittest.TestCase):
+class TestComputeGeneralizedDiceScore(unittest.TestCase):
+    @parameterized.expand([TEST_CASE_1])
+    def test_device(self, input_data, _expected_value):
+        result = compute_generalized_dice(**input_data)
+        np.testing.assert_equal(result.device, input_data["y_pred"].device)
+
     # Functional part tests
     @parameterized.expand([TEST_CASE_1, TEST_CASE_2, TEST_CASE_6, TEST_CASE_7, TEST_CASE_8, TEST_CASE_9])
     def test_value(self, input_data, expected_value):

tests/test_compute_meandice.py

@@ -191,6 +191,7 @@ class TestComputeMeanDice(unittest.TestCase):
     def test_value(self, input_data, expected_value):
         result = compute_dice(**input_data)
         np.testing.assert_allclose(result.cpu().numpy(), expected_value, atol=1e-4)
+        np.testing.assert_equal(result.device, input_data["y_pred"].device)
 
     @parameterized.expand([TEST_CASE_3])
     def test_nans(self, input_data, expected_value):

tests/test_compute_meaniou.py

@@ -191,6 +191,7 @@ class TestComputeMeanIoU(unittest.TestCase):
     def test_value(self, input_data, expected_value):
         result = compute_meaniou(**input_data)
         np.testing.assert_allclose(result.cpu().numpy(), expected_value, atol=1e-4)
+        np.testing.assert_equal(result.device, input_data["y_pred"].device)
 
     @parameterized.expand([TEST_CASE_3])
     def test_nans(self, input_data, expected_value):

tests/test_compute_panoptic_quality.py

@@ -96,6 +96,7 @@ class TestPanopticQualityMetric(unittest.TestCase):
     def test_value(self, input_params, expected_value):
         result = compute_panoptic_quality(**input_params)
         np.testing.assert_allclose(result.cpu().detach().item(), expected_value, atol=1e-4)
+        np.testing.assert_equal(result.device, input_params["pred"].device)
 
     @parameterized.expand([TEST_CLS_CASE_1, TEST_CLS_CASE_2, TEST_CLS_CASE_3, TEST_CLS_CASE_4, TEST_CLS_CASE_5])
     def test_value_class(self, input_params, y_pred, y_gt, expected_value):

tests/test_compute_variance.py

@@ -113,6 +113,7 @@ class TestComputeVariance(unittest.TestCase):
     def test_value(self, input_data, expected_value):
         result = compute_variance(**input_data)
         np.testing.assert_allclose(result.cpu().numpy(), expected_value, atol=1e-4)
+        np.testing.assert_equal(result.device, input_data["y_pred"].device)
 
     @parameterized.expand([TEST_CASE_5, TEST_CASE_6])
     def test_spatial_case(self, input_data, expected_value):

tests/test_hausdorff_distance.py

@@ -19,6 +19,8 @@
 
 from monai.metrics import HausdorffDistanceMetric
 
+_device = "cuda:0" if torch.cuda.is_available() else "cpu"
+
 
 def create_spherical_seg_3d(
     radius: float = 20.0, centre: tuple[int, int, int] = (49, 49, 49), im_shape: tuple[int, int, int] = (99, 99, 99)
@@ -116,8 +118,8 @@ def test_value(self, input_data, expected_value):
         else:
             [seg_1, seg_2] = input_data
         ct = 0
-        seg_1 = torch.tensor(seg_1)
-        seg_2 = torch.tensor(seg_2)
+        seg_1 = torch.tensor(seg_1, device=_device)
+        seg_2 = torch.tensor(seg_2, device=_device)
         for metric in ["euclidean", "chessboard", "taxicab"]:
             for directed in [True, False]:
                 hd_metric = HausdorffDistanceMetric(
@@ -130,7 +132,8 @@ def test_value(self, input_data, expected_value):
                 hd_metric(batch_seg_1, batch_seg_2)
                 result = hd_metric.aggregate(reduction="mean")
                 expected_value_curr = expected_value[ct]
-                np.testing.assert_allclose(expected_value_curr, result, rtol=1e-7)
+                np.testing.assert_allclose(expected_value_curr, result.cpu(), rtol=1e-7)
+                np.testing.assert_equal(result.device, seg_1.device)
                 ct += 1
 
     @parameterized.expand(TEST_CASES_NANS)

tests/test_label_quality_score.py

@@ -103,6 +103,7 @@ class TestLabelQualityScore(unittest.TestCase):
     def test_value(self, input_data, expected_value):
         result = label_quality_score(**input_data)
         np.testing.assert_allclose(result.cpu().numpy(), expected_value, atol=1e-4)
+        np.testing.assert_equal(result.device, input_data["y_pred"].device)
 
     @parameterized.expand([TEST_CASE_6, TEST_CASE_7])
     def test_spatial_case(self, input_data, expected_value):