Files

 master

/

fix_torch.py

Blame

Blame

Latest commit

 

History

History

24 lines (21 loc) · 766 Bytes

 master

/

fix_torch.py

Top

File metadata and controls

• Code
• Blame

24 lines (21 loc) · 766 Bytes

Raw

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

16

17

18

19

20

21

22

23

24

import importlib.util

import shutil

import os

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:

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)