update

Signed-off-by: Richard Brown <[email protected]>
Project-MONAI · Mar 17, 2021 · ad48da1 · ad48da1
1 parent 1741d80
commit ad48da1
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/monai/transforms/spatial/dictionary.py b/monai/transforms/spatial/dictionary.py
@@ -843,7 +843,8 @@ def __call__(
 
         if self._do_transform:
             cpg = self.rand_2d_elastic.deform_grid(spatial_size=sp_size)
-            cpg_w_affine, affine = self.rand_2d_elastic.rand_affine_grid(grid=cpg, return_affine=True)
+            cpg_w_affine = self.rand_2d_elastic.rand_affine_grid(grid=cpg)
+            affine = self.rand_2d_elastic.rand_affine_grid.get_transformation_matrix()
             grid = self.cpg_to_dvf(cpg_w_affine, self.rand_2d_elastic.deform_grid.spacing, sp_size)
             extra_info: Optional[Dict] = {"cpg": deepcopy(cpg), "affine": deepcopy(affine)}
         else:
@@ -1006,7 +1007,8 @@ def __call__(
             gaussian = GaussianFilter(spatial_dims=3, sigma=self.rand_3d_elastic.sigma, truncated=3.0).to(device)
             offset = torch.tensor(self.rand_3d_elastic.rand_offset, device=device).unsqueeze(0)
             grid_no_affine[:3] += gaussian(offset)[0] * self.rand_3d_elastic.magnitude
-            grid_w_affine, affine = self.rand_3d_elastic.rand_affine_grid(grid=grid_no_affine, return_affine=True)
+            grid_w_affine = self.rand_3d_elastic.rand_affine_grid(grid=grid_no_affine)
+            affine = self.rand_3d_elastic.rand_affine_grid.get_transformation_matrix()
         else:
             grid_w_affine = grid_no_affine
             affine = np.eye(len(sp_size) + 1)