fix(losses): address review comments on GlobalMutualInformationLoss buffer fix

Use register_buffer("preterm", None) / register_buffer("bin_centers", None)
unconditionally so that both buffers are always present in _buffers (with None
for b-spline). This avoids a KeyError that occurred when plain instance
attribute assignment conflicted with a subsequent register_buffer call.

Also add docstrings to the new test methods and a device-movement test that
verifies buffers follow the module when .cuda() is called.

Signed-off-by: Oleksandr Sanin <alexaaander.sanin@gmail.com>

Author: AlexanderSanin

monai/losses/image_dissimilarity.py

@@ -233,8 +233,8 @@ def __init__(
         self.kernel_type = look_up_option(kernel_type, ["gaussian", "b-spline"])
         self.num_bins = num_bins
         self.kernel_type = kernel_type
-        self.preterm: torch.Tensor
-        self.bin_centers: torch.Tensor
+        self.register_buffer("preterm", None, persistent=False)
+        self.register_buffer("bin_centers", None, persistent=False)
         if self.kernel_type == "gaussian":
             self.register_buffer("preterm", 1 / (2 * sigma**2), persistent=False)
             self.register_buffer("bin_centers", bin_centers[None, None, ...], persistent=False)

tests/losses/image_dissimilarity/test_global_mutual_information_loss.py

@@ -147,26 +147,43 @@ def test_ill_opts(self, num_bins, reduction, expected_exception, expected_messag
 
 class TestGlobalMutualInformationLossBuffers(unittest.TestCase):
     def test_gaussian_kernel_registers_buffers(self):
+        """preterm and bin_centers are registered as non-persistent buffers for gaussian kernel."""
         loss = GlobalMutualInformationLoss(kernel_type="gaussian")
-        # preterm and bin_centers must be registered buffers so .to() moves them
         self.assertIn("preterm", loss._buffers)
         self.assertIn("bin_centers", loss._buffers)
         self.assertFalse(loss.preterm.requires_grad)
         self.assertFalse(loss.bin_centers.requires_grad)
         self.assertEqual(loss.bin_centers.ndim, 3)
 
     def test_bspline_kernel_has_no_gaussian_buffers(self):
+        """b-spline kernel does not register gaussian-specific buffers."""
         loss = GlobalMutualInformationLoss(kernel_type="b-spline")
-        self.assertNotIn("preterm", loss._buffers)
-        self.assertNotIn("bin_centers", loss._buffers)
+        self.assertIsNone(loss.preterm)
+        self.assertIsNone(loss.bin_centers)
 
     def test_gaussian_kernel_forward_correct(self):
+        """Gaussian kernel forward pass returns a scalar loss."""
         pred = torch.rand(2, 1, 8, 8, dtype=torch.float32)
         target = torch.rand(2, 1, 8, 8, dtype=torch.float32)
         loss = GlobalMutualInformationLoss(kernel_type="gaussian")
         result = loss(pred, target)
         self.assertEqual(result.shape, torch.Size([]))
 
+    def test_gaussian_buffers_move_with_module(self):
+        """Buffers move to the correct device when the module is moved with .to()."""
+        loss = GlobalMutualInformationLoss(kernel_type="gaussian")
+        self.assertEqual(loss.preterm.device.type, "cpu")
+        self.assertEqual(loss.bin_centers.device.type, "cpu")
+        if not torch.cuda.is_available():
+            return
+        loss = loss.cuda()
+        self.assertEqual(loss.preterm.device.type, "cuda")
+        self.assertEqual(loss.bin_centers.device.type, "cuda")
+        pred = torch.rand(2, 1, 8, 8, device="cuda")
+        target = torch.rand(2, 1, 8, 8, device="cuda")
+        result = loss(pred, target)
+        self.assertEqual(result.device.type, "cuda")
+
 
 if __name__ == "__main__":
     unittest.main()