Ensure synchronization by adding cuda.synchronize() (#8058)

Fixes #8054

### Description
Add cuda sync after invoke cuda

### 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: YunLiu <[email protected]>

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

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: