This commit is contained in:
lalala 2026-05-09 23:53:18 +09:00 committed by GitHub
commit 69291117f6
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -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: