From e78d2304966b6265fa2320b4d87dca534ea15642 Mon Sep 17 00:00:00 2001
From: comfyanonymous <121283862+comfyanonymous@users.noreply.github.com>
Date: Wed, 23 Jul 2025 16:37:43 -0700
Subject: [PATCH] Only enable cuda malloc on cuda torch. (#9031)

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

diff --git a/cuda_malloc.py b/cuda_malloc.py
index eb2857c5f..c1d9ae3ca 100644
--- a/cuda_malloc.py
+++ b/cuda_malloc.py
@@ -74,7 +74,8 @@ if not args.cuda_malloc:
                 module = importlib.util.module_from_spec(spec)
                 spec.loader.exec_module(module)
                 version = module.__version__
-        if int(version[0]) >= 2: #enable by default for torch version 2.0 and up
+
+        if int(version[0]) >= 2 and "+cu" in version: #enable by default for torch version 2.0 and up only on cuda torch
             args.cuda_malloc = cuda_malloc_supported()
     except:
         pass