diff --git a/cuda_malloc.py b/cuda_malloc.py
index d08162cbc..b2182df37 100644
--- a/cuda_malloc.py
+++ b/cuda_malloc.py
@@ -87,7 +87,7 @@ if not args.cuda_malloc:
     except:
         pass
 
-if enables_dynamic_vram() and comfy_aimdo.control.init(0):
+if enables_dynamic_vram() and comfy_aimdo.control.init():
     args.cuda_malloc = False
     os.environ['PYTORCH_CUDA_ALLOC_CONF'] = ""