fix: use module attributes to check if cuda is available.

[lalala-233]

Current cuda check relies on "+cu" in __version__, which in torch/version.py

ComfyUI/cuda_malloc.py

Lines 66 to 82 in 9c34f5f

	version = ""
	try:
	torch_spec = importlib.util.find_spec("torch")
	for folder in torch_spec.submodule_search_locations:
	ver_file = os.path.join(folder, "version.py")
	if os.path.isfile(ver_file):
	spec = importlib.util.spec_from_file_location("torch_version_import", ver_file)
	module = importlib.util.module_from_spec(spec)
	spec.loader.exec_module(module)
	version = module.__version__
	except:
	pass
	if not args.cuda_malloc:
	try:
	if int(version[0]) >= 2 and "+cu" in version: # enable by default for torch version 2.0 and up only on cuda torch

However, some PyTorch builds don't have "+cu+" in __version__ though cuda is available. Here is my version.py

#/usr/lib/python3.14/site-packages/torch/version.py
from typing import Optional

__all__ = ['__version__', 'debug', 'cuda', 'git_version', 'hip', 'rocm', 'xpu']
__version__ = '2.11.0'
debug = False
cuda: Optional[str] = '13.2'
git_version = '70d99e998b4955e0049d13a98d77ae1b14db1f45'
hip: Optional[str] = None
rocm: Optional[str] = None
xpu: Optional[str] = None

I don't know if it is caused by Arch Linux's re-packaging or torch update its __version__ logic. But it seems better to use module attributes to directly check if cuda is available.

But I'm only test on my machine. I don't know if others' version.py are like us, so please check it before merging.

This same fix was first submitted to comfy-aimdo (see Comfy-Org/comfy-aimdo#44).

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: 6748b36e-fed6-4647-864f-f3d8393c7e83

📥 Commits

Reviewing files that changed from the base of the PR and between 9c34f5f and 9696b49.

📒 Files selected for processing (1)

• cuda_malloc.py

---

📝 Walkthrough

Walkthrough

The pull request modifies the CUDA memory allocation configuration logic in cuda_malloc.py. A new cuda_available runtime flag is introduced to track CUDA availability, initialized through dynamic detection of PyTorch's CUDA support via the torch.cuda module. The default-enable condition for CUDA memory allocation is replaced: instead of checking for "+cu" in the PyTorch version string, the code now enables by default only when PyTorch's major version is 2 or greater AND CUDA is detected as available. The existing configuration of PYTORCH_CUDA_ALLOC_CONF and environment variable handling remain unchanged. Public API signatures are preserved.

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title accurately describes the main change: replacing string-based CUDA detection with module attribute inspection.
Description check	✅ Passed	The description is directly related to the changeset, explaining the rationale for using module attributes instead of version string inspection.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

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

Tip

💬 Introducing Slack Agent: The best way for teams to turn conversations into code.

Slack Agent is built on CodeRabbit's deep understanding of your code, so your team can collaborate across the entire SDLC without losing context.

• Generate code and open pull requests
• Plan features and break down work
• Investigate incidents and troubleshoot customer tickets together
• Automate recurring tasks and respond to alerts with triggers
• Summarize progress and report instantly

Built for teams:

• Shared memory across your entire org—no repeating context
• Per-thread sandboxes to safely plan and execute work
• Governance built-in—scoped access, auditability, and budget controls

One agent for your entire SDLC. Right inside Slack.

👉 Get started

---

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

• X
• Mastodon
• Reddit
• LinkedIn

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