From 9696b494fac9e51dea8003d7f73ff85da37596ff Mon Sep 17 00:00:00 2001
From: lalala_233 <lalala_233@qq.com>
Date: Wed, 6 May 2026 22:59:15 +0800
Subject: [PATCH] fix: use module attributes to check if cuda is available.

---
 cuda_malloc.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/cuda_malloc.py b/cuda_malloc.py
index f7651981c..722bc2b20 100644
--- a/cuda_malloc.py
+++ b/cuda_malloc.py
@@ -64,6 +64,7 @@ def cuda_malloc_supported():
 
 
 version = ""
+cuda_available = False
 
 try:
     torch_spec = importlib.util.find_spec("torch")
@@ -74,12 +75,13 @@ try:
             module = importlib.util.module_from_spec(spec)
             spec.loader.exec_module(module)
             version = module.__version__
+            cuda_available = module.cuda != None
 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
+        if int(version[0]) >= 2 and cuda_available:  # enable by default for torch version 2.0 and up only on cuda torch
             if PerformanceFeature.AutoTune not in args.fast:  # Autotune has issues with cuda malloc
                 args.cuda_malloc = cuda_malloc_supported()
     except: