From e6ec5899fd156b57e58cec5f070aa30e25c16704 Mon Sep 17 00:00:00 2001 From: Philip Yang Date: Sat, 20 Oct 2018 20:39:44 +0000 Subject: [PATCH] Find detectron ops from pytorch nightly build --- detectron/utils/env.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/detectron/utils/env.py b/detectron/utils/env.py index ea66faacd..7e8b25ed9 100644 --- a/detectron/utils/env.py +++ b/detectron/utils/env.py @@ -22,6 +22,7 @@ import os import sys +import torch # Default value of the CMake install prefix _CMAKE_INSTALL_PREFIX = '/usr/local' @@ -61,7 +62,7 @@ def import_nccl_ops(): def get_detectron_ops_lib(): """Retrieve Detectron ops library.""" # Candidate prefixes for detectron ops lib path - prefixes = [_CMAKE_INSTALL_PREFIX, sys.prefix, sys.exec_prefix] + sys.path + prefixes = [_CMAKE_INSTALL_PREFIX, sys.prefix, sys.exec_prefix, os.path.dirname(torch.__file__)] + sys.path # Candidate subdirs for detectron ops lib subdirs = ['lib', 'torch/lib'] # Try to find detectron ops lib