diff --git a/monai/networks/layers/filtering.py b/monai/networks/layers/filtering.py index 0ff1187dcc..c48c77cf98 100644 --- a/monai/networks/layers/filtering.py +++ b/monai/networks/layers/filtering.py @@ -51,6 +51,8 @@ def forward(ctx, input, spatial_sigma=5, color_sigma=0.5, fast_approx=True): ctx.cs = color_sigma ctx.fa = fast_approx output_data = _C.bilateral_filter(input, spatial_sigma, color_sigma, fast_approx) + if torch.cuda.is_available(): + torch.cuda.synchronize() return output_data @staticmethod @@ -139,7 +141,8 @@ def forward(ctx, input_img, sigma_x, sigma_y, sigma_z, color_sigma): do_dsig_y, do_dsig_z, ) - + if torch.cuda.is_available(): + torch.cuda.synchronize() return output_tensor @staticmethod @@ -301,7 +304,8 @@ def forward(ctx, input_img, guidance_img, sigma_x, sigma_y, sigma_z, color_sigma do_dsig_z, guidance_img, ) - + if torch.cuda.is_available(): + torch.cuda.synchronize() return output_tensor @staticmethod diff --git a/monai/transforms/utils.py b/monai/transforms/utils.py index 7027c07d67..1d1f070568 100644 --- a/monai/transforms/utils.py +++ b/monai/transforms/utils.py @@ -2512,6 +2512,7 @@ def distance_transform_edt( block_params=block_params, float64_distances=float64_distances, ) + torch.cuda.synchronize() else: if not has_ndimage: raise RuntimeError("scipy.ndimage required if cupy is not available") @@ -2545,7 +2546,7 @@ def distance_transform_edt( r_vals = [] if return_distances and distances_original is None: - r_vals.append(distances) + r_vals.append(distances_ if use_cp else distances) if return_indices and indices_original is None: r_vals.append(indices) if not r_vals: