Skip to content

Conversation

@benediktjohannes
Copy link

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

…abel 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 <benedikt.johannes.hofer@gmail.com>
@coderabbitai
Copy link
Contributor

coderabbitai bot commented Jan 17, 2026

📝 Walkthrough

Walkthrough

Refactors in monai/auto3dseg/analyzer.py change per-sample data access to use local tensors (image_tensor, label_tensor) and index them directly. CUDA detection now checks both tensors with any(...). Per-sample foreground placeholders switch from torch.Tensor([0]) to MetaTensor([0.0]). Shape checks and overall statistic computation/update flow are unchanged. No public API or signature changes.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~10 minutes

🚥 Pre-merge checks | ✅ 2 | ❌ 1
❌ Failed checks (1 inconclusive)
Check name Status Explanation Resolution
Description check ❓ Inconclusive PR description explains the changes and their purpose but doesn't follow the repository's template structure with required sections like Types of changes checkboxes. Add the Types of changes section with appropriate checkboxes marked as specified in the repository template.
✅ Passed checks (2 passed)
Check name Status Explanation
Title check ✅ Passed Title clearly and concisely summarizes the main change: fixing a GPU memory leak by checking both image and label tensors for CUDA device detection.
Docstring Coverage ✅ Passed Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.

✏️ Tip: You can configure your own custom pre-merge checks in the settings.

✨ Finishing touches
  • 📝 Generate docstrings

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

Comment @coderabbitai help to get the list of available commands and usage tips.

Copy link
Contributor

@coderabbitai coderabbitai bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 1

🤖 Fix all issues with AI agents
In `@monai/auto3dseg/analyzer.py`:
- Around line 471-476: The code uses undefined image_tensor and label_tensor
before they are assigned; extract them from the incoming data dict like
ImageStats and FgImageStats do (use the same keys those classes expect, e.g.,
"image" for the input image and "label" for the segmentation label), then
proceed with the existing logic (compute using_cuda, save restore_grad_state,
set_grad_enabled(False), build ndas from image_tensor and cast label_tensor to
int16 for ndas_label). Ensure the variables are set to the tensor objects (with
correct shape [B,...]) before the lines referencing using_cuda, ndas, and
ndas_label.

Signed-off-by: benediktjohannes <benedikt.johannes.hofer@gmail.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant