add -lcurand to fix torch-nightly issue w. JIT (#1688)

This commit is contained in:
Jeff Rasley 2022-01-10 13:47:56 -08:00 коммит произвёл GitHub
Родитель 3a4cb04243
Коммит a85dce0728
Не найден ключ, соответствующий данной подписи
Идентификатор ключа GPG: 4AEE18F83AFDEB23
1 изменённых файлов: 3 добавлений и 0 удалений

Просмотреть файл

@ -556,6 +556,9 @@ class CUDAOpBuilder(OpBuilder):
class TorchCPUOpBuilder(CUDAOpBuilder):
def extra_ldflags(self):
return ['-lcurand']
def cxx_args(self):
import torch
CUDA_LIB64 = os.path.join(torch.utils.cpp_extension.CUDA_HOME, "lib64")