Fix non-functional jitter in Warp.get_reference_grid

monai/networks/blocks/warp.py

@@ -121,12 +121,13 @@ def get_reference_grid(self, ddf: torch.Tensor, jitter: bool = False, seed: int
         mesh_points = [torch.arange(0, dim) for dim in ddf.shape[2:]]
         grid = torch.stack(meshgrid_ij(*mesh_points), dim=0)  # (spatial_dims, ...)
         grid = torch.stack([grid] * ddf.shape[0], dim=0)  # (batch, spatial_dims, ...)
-        self.ref_grid = grid.to(ddf)
+        grid = grid.to(ddf)
         if jitter:
             # Define reference grid on non-integer values
-            with torch.random.fork_rng(enabled=seed):
+            with torch.random.fork_rng():
                 torch.random.manual_seed(seed)
                 grid += torch.rand_like(grid)
+        self.ref_grid = grid
         self.ref_grid.requires_grad = False
         return self.ref_grid
 

tests/networks/blocks/warp/test_warp.py

@@ -138,6 +138,21 @@ def test_ill_shape(self):
         with self.assertRaisesRegex(ValueError, ""):
             warp_layer(image=torch.arange(4).reshape((1, 1, 2, 2)).to(dtype=torch.float), ddf=torch.zeros(1, 2, 3, 3))
 
+    def test_jitter(self):
+        ddf = torch.zeros(1, 2, 4, 5)
+        grid = Warp(jitter=True).get_reference_grid(ddf, jitter=True, seed=0)
+        self.assertTrue(grid.is_floating_point())
+        self.assertTrue(bool((grid != grid.round()).any()))
+
+        ref = Warp().get_reference_grid(ddf, jitter=False)
+        self.assertTrue(bool((ref == ref.round()).all()))
+
+        same = Warp().get_reference_grid(ddf, jitter=True, seed=7)
+        repeat = Warp().get_reference_grid(ddf, jitter=True, seed=7)
+        other = Warp().get_reference_grid(ddf, jitter=True, seed=8)
+        self.assertTrue(torch.equal(same, repeat))
+        self.assertFalse(torch.equal(same, other))
+
     def test_grad(self):
         for b in GridSampleMode:
             for p in GridSamplePadMode: