diff --git a/fix_torch.py b/fix_torch.py index b60a0bbb..e350f5c7 100644 --- a/fix_torch.py +++ b/fix_torch.py @@ -9,8 +9,16 @@ 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: + 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"), os.path.join(lib_folder, "libomp140.x86_64.dll")) + shutil.copyfile(os.path.join(lib_folder, "libiomp5md.dll"), dest)