From aa93e6c3a7873f79583897c9074ba522949d81d2 Mon Sep 17 00:00:00 2001 From: benediktjohannes Date: Sat, 17 Jan 2026 16:28:57 +0100 Subject: [PATCH 1/2] Fixes #8697 GPU memory leak by checking both image and label tensors for CUDA device Modified device detection to check BOTH image and label tensors torch.cuda.empty_cache() now called if EITHER tensor is on GPU Prevents GPU memory leaks in mixed device scenarios Signed-off-by: benediktjohannes --- monai/auto3dseg/analyzer.py | 11 ++++------- 1 file changed, 4 insertions(+), 7 deletions(-) diff --git a/monai/auto3dseg/analyzer.py b/monai/auto3dseg/analyzer.py index 8d662df83d..a2d8c77e25 100644 --- a/monai/auto3dseg/analyzer.py +++ b/monai/auto3dseg/analyzer.py @@ -468,21 +468,18 @@ 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 + using_cuda = any(isinstance(t, torch.Tensor) 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)): From f8cc9fff4268ad87a456fb4fe27483d5095e3ed3 Mon Sep 17 00:00:00 2001 From: benediktjohannes Date: Sat, 17 Jan 2026 16:47:10 +0100 Subject: [PATCH 2/2] Forgot to define variables, thanks coderabbitai for the information Signed-off-by: benediktjohannes --- monai/auto3dseg/analyzer.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/monai/auto3dseg/analyzer.py b/monai/auto3dseg/analyzer.py index a2d8c77e25..736607a815 100644 --- a/monai/auto3dseg/analyzer.py +++ b/monai/auto3dseg/analyzer.py @@ -468,6 +468,8 @@ def __call__(self, data: Mapping[Hashable, MetaTensor]) -> dict[Hashable, MetaTe """ d: dict[Hashable, MetaTensor] = dict(data) start = time.time() + image_tensor = d[self.image_key] + label_tensor = d[self.label_key] using_cuda = any(isinstance(t, torch.Tensor) and t.device.type == "cuda" for t in (image_tensor, label_tensor)) restore_grad_state = torch.is_grad_enabled() torch.set_grad_enabled(False)