Skip to content

GitLab

Unverified Commit 3b83ef51 authored by msbaines's avatar msbaines Committed by GitHub
Browse files

[fix] make extension build robust to include path (#196) 

Fixes #190
parent cd496b36
Hide whitespace changes
Inline Side-by-side
Showing with 3 additions and 0 deletions
setup.py
View file @ 3b83ef51
...@@ -10,6 +10,8 @@ import setuptools ...@@ -10,6 +10,8 @@ import setuptools
import torch import torch
from torch.utils.cpp_extension import CUDA_HOME, BuildExtension, CUDAExtension from torch.utils.cpp_extension import CUDA_HOME, BuildExtension, CUDAExtension
this_dir = os.path.dirname(os.path.abspath(__file__))
def fetch_requirements(): def fetch_requirements():
with open("requirements.txt") as f: with open("requirements.txt") as f:
...@@ -35,6 +37,7 @@ if (torch.cuda.is_available() and CUDA_HOME is not None) or force_cuda: ...@@ -35,6 +37,7 @@ if (torch.cuda.is_available() and CUDA_HOME is not None) or force_cuda:
[ [
CUDAExtension( CUDAExtension(
name="fairscale.fused_adam_cuda", name="fairscale.fused_adam_cuda",
include_dirs=[os.path.join(this_dir, "fairscale/clib/fused_adam_cuda")],
sources=[ sources=[
"fairscale/clib/fused_adam_cuda/fused_adam_cuda.cpp", "fairscale/clib/fused_adam_cuda/fused_adam_cuda.cpp",
"fairscale/clib/fused_adam_cuda/fused_adam_cuda_kernel.cu", "fairscale/clib/fused_adam_cuda/fused_adam_cuda_kernel.cu",
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment