From 73b26e53754ffc9842941ce6b763668de5653829 Mon Sep 17 00:00:00 2001
From: huchenlei <huchenlei@proton.me>
Date: Sun, 8 Dec 2024 21:01:41 -0500
Subject: [PATCH] nit

---
 fix_torch.py | 36 ++++++++++++++++++++----------------
 main.py      |  4 ++--
 2 files changed, 22 insertions(+), 18 deletions(-)

diff --git a/fix_torch.py b/fix_torch.py
index e350f5c7d..4aecb23f0 100644
--- a/fix_torch.py
+++ b/fix_torch.py
@@ -5,20 +5,24 @@ import ctypes
 import logging
 
 
-torch_spec = importlib.util.find_spec("torch")
-for folder in torch_spec.submodule_search_locations:
-    lib_folder = os.path.join(folder, "lib")
-    test_file = os.path.join(lib_folder, "fbgemm.dll")
-    dest = os.path.join(lib_folder, "libomp140.x86_64.dll")
-    if os.path.exists(dest):
-        break
-
-    with open(test_file, 'rb') as f:
-        contents = f.read()
-        if b"libomp140.x86_64.dll" not in contents:
+def fix_pytorch_libomp():
+    """
+    Fix PyTorch libomp DLL issue on Windows by copying the correct DLL file if needed.
+    """
+    torch_spec = importlib.util.find_spec("torch")
+    for folder in torch_spec.submodule_search_locations:
+        lib_folder = os.path.join(folder, "lib")
+        test_file = os.path.join(lib_folder, "fbgemm.dll")
+        dest = os.path.join(lib_folder, "libomp140.x86_64.dll")
+        if os.path.exists(dest):
             break
-    try:
-        mydll = ctypes.cdll.LoadLibrary(test_file)
-    except FileNotFoundError as e:
-        logging.warning("Detected pytorch version with libomp issue, patching.")
-        shutil.copyfile(os.path.join(lib_folder, "libiomp5md.dll"), dest)
+
+        with open(test_file, "rb") as f:
+            contents = f.read()
+            if b"libomp140.x86_64.dll" not in contents:
+                break
+        try:
+            mydll = ctypes.cdll.LoadLibrary(test_file)
+        except FileNotFoundError as e:
+            logging.warning("Detected pytorch version with libomp issue, patching.")
+            shutil.copyfile(os.path.join(lib_folder, "libiomp5md.dll"), dest)
diff --git a/main.py b/main.py
index ceaa9d809..87ce1526b 100644
--- a/main.py
+++ b/main.py
@@ -86,9 +86,9 @@ if __name__ == "__main__":
     import cuda_malloc
 
 if args.windows_standalone_build:
-    # TODO: Convert fix_torch to a function.
     try:
-        import fix_torch # noqa: F401
+        from fix_torch import fix_pytorch_libomp
+        fix_pytorch_libomp()
     except:
         pass