Fixes #8697 GPU memory leak by checking both image and label tensors for CUDA device

monai/apps/detection/transforms/box_ops.py

@@ -267,7 +267,7 @@ def convert_box_to_mask(
             boxes_only_mask = np.ones(box_size, dtype=np.int16) * np.int16(labels_np[b])
         # apply to global mask
         slicing = [b]
-        slicing.extend(slice(boxes_np[b, d], boxes_np[b, d + spatial_dims]) for d in range(spatial_dims))  # type:ignore
+        slicing.extend(slice(boxes_np[b, d], boxes_np[b, d + spatial_dims]) for d in range(spatial_dims))  # type: ignore
         boxes_mask_np[tuple(slicing)] = boxes_only_mask
     return convert_to_dst_type(src=boxes_mask_np, dst=boxes, dtype=torch.int16)[0]
 

monai/apps/detection/transforms/dictionary.py

@@ -125,10 +125,8 @@ def __init__(self, box_keys: KeysCollection, box_ref_image_keys: str, allow_miss
         super().__init__(box_keys, allow_missing_keys)
         box_ref_image_keys_tuple = ensure_tuple(box_ref_image_keys)
         if len(box_ref_image_keys_tuple) > 1:
-            raise ValueError(
-                "Please provide a single key for box_ref_image_keys.\
-                All boxes of box_keys are attached to box_ref_image_keys."
-            )
+            raise ValueError("Please provide a single key for box_ref_image_keys.\
+                All boxes of box_keys are attached to box_ref_image_keys.")
         self.box_ref_image_keys = box_ref_image_keys
 
     def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> dict[Hashable, NdarrayOrTensor]:
@@ -289,10 +287,8 @@ def __init__(
         super().__init__(box_keys, allow_missing_keys)
         box_ref_image_keys_tuple = ensure_tuple(box_ref_image_keys)
         if len(box_ref_image_keys_tuple) > 1:
-            raise ValueError(
-                "Please provide a single key for box_ref_image_keys.\
-                All boxes of box_keys are attached to box_ref_image_keys."
-            )
+            raise ValueError("Please provide a single key for box_ref_image_keys.\
+                All boxes of box_keys are attached to box_ref_image_keys.")
         self.box_ref_image_keys = box_ref_image_keys
         self.image_meta_key = image_meta_key or f"{box_ref_image_keys}_{image_meta_key_postfix}"
         self.converter_to_image_coordinate = AffineBox()
@@ -310,10 +306,8 @@ def extract_affine(self, data: Mapping[Hashable, torch.Tensor]) -> tuple[Ndarray
         else:
             raise ValueError(f"{meta_key} is not found. Please check whether it is the correct the image meta key.")
         if "affine" not in meta_dict:
-            raise ValueError(
-                f"'affine' is not found in {meta_key}. \
-                Please check whether it is the correct the image meta key."
-            )
+            raise ValueError(f"'affine' is not found in {meta_key}. \
+                Please check whether it is the correct the image meta key.")
         affine: NdarrayOrTensor = meta_dict["affine"]
 
         if self.affine_lps_to_ras:  # RAS affine
@@ -815,16 +809,12 @@ def __init__(
     ) -> None:
         box_keys_tuple = ensure_tuple(box_keys)
         if len(box_keys_tuple) != 1:
-            raise ValueError(
-                "Please provide a single key for box_keys.\
-                All label_keys are attached to this box_keys."
-            )
+            raise ValueError("Please provide a single key for box_keys.\
+                All label_keys are attached to this box_keys.")
         box_ref_image_keys_tuple = ensure_tuple(box_ref_image_keys)
         if len(box_ref_image_keys_tuple) != 1:
-            raise ValueError(
-                "Please provide a single key for box_ref_image_keys.\
-                All box_keys and label_keys are attached to this box_ref_image_keys."
-            )
+            raise ValueError("Please provide a single key for box_ref_image_keys.\
+                All box_keys and label_keys are attached to this box_ref_image_keys.")
         self.label_keys = ensure_tuple(label_keys)
         super().__init__(box_keys_tuple, allow_missing_keys)
 
@@ -1091,10 +1081,8 @@ def __init__(
 
         box_keys_tuple = ensure_tuple(box_keys)
         if len(box_keys_tuple) != 1:
-            raise ValueError(
-                "Please provide a single key for box_keys.\
-                All label_keys are attached to this box_keys."
-            )
+            raise ValueError("Please provide a single key for box_keys.\
+                All label_keys are attached to this box_keys.")
         self.box_keys = box_keys_tuple[0]
         self.label_keys = ensure_tuple(label_keys)
 

monai/apps/detection/utils/anchor_utils.py

@@ -124,10 +124,8 @@ def __init__(
             aspect_ratios = (aspect_ratios,) * len(self.sizes)
 
         if len(self.sizes) != len(aspect_ratios):
-            raise ValueError(
-                "len(sizes) and len(aspect_ratios) should be equal. \
-                It represents the number of feature maps."
-            )
+            raise ValueError("len(sizes) and len(aspect_ratios) should be equal. \
+                It represents the number of feature maps.")
 
         spatial_dims = len(ensure_tuple(aspect_ratios[0][0])) + 1
         spatial_dims = look_up_option(spatial_dims, [2, 3])
@@ -172,16 +170,12 @@ def generate_anchors(
         scales_t = torch.as_tensor(scales, dtype=dtype, device=device)  # sized (N,)
         aspect_ratios_t = torch.as_tensor(aspect_ratios, dtype=dtype, device=device)  # sized (M,) or (M,2)
         if (self.spatial_dims >= 3) and (len(aspect_ratios_t.shape) != 2):
-            raise ValueError(
-                f"In {self.spatial_dims}-D image, aspect_ratios for each level should be \
-                {len(aspect_ratios_t.shape) - 1}-D. But got aspect_ratios with shape {aspect_ratios_t.shape}."
-            )
+            raise ValueError(f"In {self.spatial_dims}-D image, aspect_ratios for each level should be \
+                {len(aspect_ratios_t.shape) - 1}-D. But got aspect_ratios with shape {aspect_ratios_t.shape}.")
 
         if (self.spatial_dims >= 3) and (aspect_ratios_t.shape[1] != self.spatial_dims - 1):
-            raise ValueError(
-                f"In {self.spatial_dims}-D image, aspect_ratios for each level should has \
-                shape (_,{self.spatial_dims - 1}). But got aspect_ratios with shape {aspect_ratios_t.shape}."
-            )
+            raise ValueError(f"In {self.spatial_dims}-D image, aspect_ratios for each level should has \
+                shape (_,{self.spatial_dims - 1}). But got aspect_ratios with shape {aspect_ratios_t.shape}.")
 
         # if 2d, w:h = 1:aspect_ratios
         if self.spatial_dims == 2:

monai/apps/reconstruction/transforms/array.py

@@ -61,10 +61,8 @@ def __init__(
                 real/imaginary parts.
         """
         if len(center_fractions) != len(accelerations):
-            raise ValueError(
-                "Number of center fractions \
-                should match number of accelerations"
-            )
+            raise ValueError("Number of center fractions \
+                should match number of accelerations")
 
         self.center_fractions = center_fractions
         self.accelerations = accelerations

monai/auto3dseg/analyzer.py

@@ -105,7 +105,7 @@ def update_ops_nested_label(self, nested_key: str, op: Operations) -> None:
             raise ValueError("Nested_key input format is wrong. Please ensure it is like key1#0#key2")
         root: str
         child_key: str
-        (root, _, child_key) = keys
+        root, _, child_key = keys
         if root not in self.ops:
             self.ops[root] = [{}]
         self.ops[root][0].update({child_key: None})
@@ -468,21 +468,22 @@ def __call__(self, data: Mapping[Hashable, MetaTensor]) -> dict[Hashable, MetaTe
         """
         d: dict[Hashable, MetaTensor] = dict(data)
         start = time.time()
-        if isinstance(d[self.image_key], (torch.Tensor, MetaTensor)) and d[self.image_key].device.type == "cuda":
-            using_cuda = True
-        else:
-            using_cuda = False
+        image_tensor = d[self.image_key]
+        label_tensor = d[self.label_key]
+        using_cuda = any(
+            isinstance(t, (torch.Tensor, MetaTensor)) and t.device.type == "cuda" for t in (image_tensor, label_tensor)
+        )
         restore_grad_state = torch.is_grad_enabled()
         torch.set_grad_enabled(False)
 
-        ndas: list[MetaTensor] = [d[self.image_key][i] for i in range(d[self.image_key].shape[0])]  # type: ignore
-        ndas_label: MetaTensor = d[self.label_key].astype(torch.int16)  # (H,W,D)
+        ndas: list[MetaTensor] = [image_tensor[i] for i in range(image_tensor.shape[0])]  # type: ignore
+        ndas_label: MetaTensor = label_tensor.astype(torch.int16)  # (H,W,D)
 
         if ndas_label.shape != ndas[0].shape:
             raise ValueError(f"Label shape {ndas_label.shape} is different from image shape {ndas[0].shape}")
 
         nda_foregrounds: list[torch.Tensor] = [get_foreground_label(nda, ndas_label) for nda in ndas]
-        nda_foregrounds = [nda if nda.numel() > 0 else torch.Tensor([0]) for nda in nda_foregrounds]
+        nda_foregrounds = [nda if nda.numel() > 0 else MetaTensor([0.0]) for nda in nda_foregrounds]
 
         unique_label = unique(ndas_label)
         if isinstance(ndas_label, (MetaTensor, torch.Tensor)):

monai/bundle/scripts.py

@@ -1948,7 +1948,7 @@ def create_workflow(
 
     """
     _args = update_kwargs(args=args_file, workflow_name=workflow_name, config_file=config_file, **kwargs)
-    (workflow_name, config_file) = _pop_args(
+    workflow_name, config_file = _pop_args(
         _args, workflow_name=ConfigWorkflow, config_file=None
     )  # the default workflow name is "ConfigWorkflow"
     if isinstance(workflow_name, str):

monai/bundle/utils.py

@@ -124,10 +124,8 @@
         "run_name": None,
         # may fill it at runtime
         "save_execute_config": True,
-        "is_not_rank0": (
-            "$torch.distributed.is_available() \
-                and torch.distributed.is_initialized() and torch.distributed.get_rank() > 0"
-        ),
+        "is_not_rank0": ("$torch.distributed.is_available() \
+                and torch.distributed.is_initialized() and torch.distributed.get_rank() > 0"),
         # MLFlowHandler config for the trainer
         "trainer": {
             "_target_": "MLFlowHandler",

monai/data/dataset.py

@@ -139,7 +139,7 @@ class DatasetFunc(Dataset):
     """
 
     def __init__(self, data: Any, func: Callable, **kwargs) -> None:
-        super().__init__(data=None, transform=None)  # type:ignore
+        super().__init__(data=None, transform=None)  # type: ignore
         self.src = data
         self.func = func
         self.kwargs = kwargs
@@ -1635,7 +1635,7 @@ def _cachecheck(self, item_transformed):
                         return (_data, _meta)
                     return _data
                 else:
-                    item: list[dict[Any, Any]] = [{} for _ in range(len(item_transformed))]  # type:ignore
+                    item: list[dict[Any, Any]] = [{} for _ in range(len(item_transformed))]  # type: ignore
                     for i, _item in enumerate(item_transformed):
                         for k in _item:
                             meta_i_k = self._load_meta_cache(meta_hash_file_name=f"{hashfile.name}-{k}-meta-{i}")

monai/handlers/utils.py

@@ -48,7 +48,7 @@ def stopping_fn_from_loss() -> Callable[[Engine], Any]:
     """
 
     def stopping_fn(engine: Engine) -> Any:
-        return -engine.state.output  # type:ignore
+        return -engine.state.output  # type: ignore
 
     return stopping_fn
 

monai/losses/dice.py

@@ -203,11 +203,9 @@ def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
                 self.class_weight = torch.as_tensor([self.class_weight] * num_of_classes)
             else:
                 if self.class_weight.shape[0] != num_of_classes:
-                    raise ValueError(
-                        """the length of the `weight` sequence should be the same as the number of classes.
+                    raise ValueError("""the length of the `weight` sequence should be the same as the number of classes.
                         If `include_background=False`, the weight should not include
-                        the background category class 0."""
-                    )
+                        the background category class 0.""")
             if self.class_weight.min() < 0:
                 raise ValueError("the value/values of the `weight` should be no less than 0.")
             # apply class_weight to loss

monai/losses/focal_loss.py

@@ -183,11 +183,9 @@ def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
                 self.class_weight = torch.as_tensor([self.class_weight] * num_of_classes)
             else:
                 if self.class_weight.shape[0] != num_of_classes:
-                    raise ValueError(
-                        """the length of the `weight` sequence should be the same as the number of classes.
+                    raise ValueError("""the length of the `weight` sequence should be the same as the number of classes.
                         If `include_background=False`, the weight should not include
-                        the background category class 0."""
-                    )
+                        the background category class 0.""")
             if self.class_weight.min() < 0:
                 raise ValueError("the value/values of the `weight` should be no less than 0.")
             # apply class_weight to loss

monai/metrics/utils.py

@@ -320,7 +320,7 @@ def get_edge_surface_distance(
     edges_spacing = None
     if use_subvoxels:
         edges_spacing = spacing if spacing is not None else ([1] * len(y_pred.shape))
-    (edges_pred, edges_gt, *areas) = get_mask_edges(
+    edges_pred, edges_gt, *areas = get_mask_edges(
         y_pred, y, crop=True, spacing=edges_spacing, always_return_as_numpy=False
     )
     if not edges_gt.any():

monai/transforms/io/array.py

@@ -11,6 +11,7 @@
 """
 A collection of "vanilla" transforms for IO functions.
 """
+
 from __future__ import annotations
 
 import inspect

monai/transforms/utility/array.py

@@ -702,7 +702,7 @@ def __init__(
                 # if the root log level is higher than INFO, set a separate stream handler to record
                 console = logging.StreamHandler(sys.stdout)
                 console.setLevel(logging.INFO)
-                console.is_data_stats_handler = True  # type:ignore[attr-defined]
+                console.is_data_stats_handler = True  # type: ignore[attr-defined]
                 _logger.addHandler(console)
 
     def __call__(

tests/integration/test_loader_semaphore.py

@@ -10,6 +10,7 @@
 # limitations under the License.
 """this test should not generate errors or
 UserWarning: semaphore_tracker: There appear to be 1 leaked semaphores"""
+
 from __future__ import annotations
 
 import multiprocessing as mp

tests/profile_subclass/profiling.py

@@ -12,6 +12,7 @@
 Comparing torch.Tensor, SubTensor, SubWithTorchFunc, MetaTensor
 Adapted from https://github.com/pytorch/pytorch/tree/v1.11.0/benchmarks/overrides_benchmark
 """
+
 from __future__ import annotations
 
 import argparse

tests/profile_subclass/pyspy_profiling.py

@@ -12,6 +12,7 @@
 To be used with py-spy, comparing torch.Tensor, SubTensor, SubWithTorchFunc, MetaTensor
 Adapted from https://github.com/pytorch/pytorch/tree/v1.11.0/benchmarks/overrides_benchmark
 """
+
 from __future__ import annotations
 
 import argparse

versioneer.py

@@ -273,6 +273,7 @@
 [travis-url]: https://travis-ci.com/github/python-versioneer/python-versioneer
 
 """
+
 # pylint:disable=invalid-name,import-outside-toplevel,missing-function-docstring
 # pylint:disable=missing-class-docstring,too-many-branches,too-many-statements
 # pylint:disable=raise-missing-from,too-many-lines,too-many-locals,import-error
@@ -428,9 +429,7 @@ def run_command(commands, args, cwd=None, verbose=False, hide_stderr=False, env=
     return stdout, process.returncode
 
 
-LONG_VERSION_PY[
-    "git"
-] = r'''
+LONG_VERSION_PY["git"] = r'''
 # This file helps to compute a version number in source trees obtained from
 # git-archive tarball (such as those provided by githubs download-from-tag
 # feature). Distribution tarballs (built by setup.py sdist) and build