pytorch · dagitses · Mar 28, 2022 · Mar 28, 2022 · Mar 28, 2022 · Mar 28, 2022
diff --git a/tools/setup_helpers/generate_code.py b/tools/setup_helpers/generate_code.py
@@ -46,10 +46,8 @@ def generate_code(ninja_global: Optional[str] = None,
     else:
         python_install_dir = install_dir
     autograd_gen_dir = os.path.join(install_dir, 'autograd', 'generated')
-    jit_gen_dir = os.path.join(install_dir, 'jit', 'generated')
-    for d in (autograd_gen_dir, jit_gen_dir, python_install_dir):
-        if not os.path.exists(d):
-            os.makedirs(d)
+    for d in (autograd_gen_dir, python_install_dir):
+        os.makedirs(d, exist_ok=True)
     autograd_dir = os.fspath(pathlib.Path(__file__).parent.parent / "autograd")
 
     if subset == "pybindings" or not subset:
@@ -189,8 +187,7 @@ def main() -> None:
         if options.install_dir is None:
             options.install_dir = "torch/csrc"
         lazy_install_dir = os.path.join(options.install_dir, "lazy/generated")
-        if not os.path.exists(lazy_install_dir):
-            os.makedirs(lazy_install_dir)
+        os.makedirs(lazy_install_dir, exist_ok=True)
 
         assert os.path.isfile(ts_backend_yaml), f"Unable to access ts_backend_yaml: {ts_backend_yaml}"
         assert os.path.isfile(ts_native_functions), f"Unable to access {ts_native_functions}"