Update setup.py (#269)
parent
27024e42da
commit
2b62f419c2
12
setup.py
12
setup.py
|
@ -24,6 +24,18 @@ import glob
|
|||
import os
|
||||
import subprocess
|
||||
|
||||
import subprocess
|
||||
import sys
|
||||
|
||||
def install_torch():
|
||||
try:
|
||||
import torch
|
||||
except ImportError:
|
||||
subprocess.check_call([sys.executable, "-m", "pip", "install", "torch"])
|
||||
|
||||
# Call the function to ensure torch is installed
|
||||
install_torch()
|
||||
|
||||
import torch
|
||||
from setuptools import find_packages, setup
|
||||
from torch.utils.cpp_extension import CUDA_HOME, CppExtension, CUDAExtension
|
||||
|
|
Loading…
Reference in New Issue