v4.3 tag release update. (#2789)

examples/python/CuTeDSL/cute/ffi/jit_argument.py

@@ -243,12 +243,7 @@ import tempfile
 import torch
 
 
-def run_test(tmpdir=None, cmake_args=""):
-    # Skip cleanup if user provides tmpdir
-    cleanup = tmpdir is None
-    # Initialize temporary build directory
-    tmpdir = tmpdir or tempfile.mkdtemp()
-
+def run_test(tmpdir=None, cmake_args="", cleanup=True):
     try:
         current_dir = os.path.dirname(os.path.abspath(__file__))
 
@@ -256,8 +251,6 @@ def run_test(tmpdir=None, cmake_args=""):
         subprocess.run(["cmake", "-B", tmpdir, current_dir] + cmake_args, check=True)
         subprocess.run(["cmake", "--build", tmpdir], check=True)
 
-        sys.path.append(tmpdir)
-
         from tensor import make_tensor, pycapsule_get_pointer
 
         # Mock test tensor and corresponding C structure for this example
@@ -314,4 +307,13 @@ if __name__ == "__main__":
     )
     args = parser.parse_args()
 
-    run_test(tmpdir=args.tmp_dir, cmake_args=args.cmake_args)
+    if args.tmp_dir:
+        tmp_dir = args.tmp_dir
+        cleanup = False
+    else:
+        tmp_dir = tempfile.mkdtemp()
+        cleanup = True
+
+    sys.path.append(tmp_dir)
+
+    run_test(tmp_dir, args.cmake_args, cleanup)