Find detectron ops lib in pytorch nightly anaconda package

Reviewed By: rbgirshick

Differential Revision: D10433948

fbshipit-source-id: 0d4dbc7f532e13c559da65c403aa06238f612e25

detectron/utils/env.py

@@ -25,6 +25,8 @@ import sys
 
 # Default value of the CMake install prefix
 _CMAKE_INSTALL_PREFIX = '/usr/local'
+# Detectron ops lib
+_DETECTRON_OPS_LIB = 'libcaffe2_detectron_ops_gpu.so'
 
 
 def get_runtime_dir():
@@ -58,19 +60,18 @@ def import_nccl_ops():
 
 def get_detectron_ops_lib():
     """Retrieve Detectron ops library."""
-    # Candidate prefixes for the detectron ops lib path
+    # Candidate prefixes for detectron ops lib path
     prefixes = [_CMAKE_INSTALL_PREFIX, sys.prefix, sys.exec_prefix] + sys.path
-    # Search for detectron ops lib
+    # Candidate subdirs for detectron ops lib
+    subdirs = ['lib', 'torch/lib']
+    # Try to find detectron ops lib
     for prefix in prefixes:
-        ops_path = os.path.join(prefix, 'lib/libcaffe2_detectron_ops_gpu.so')
-        if os.path.exists(ops_path):
-            # TODO(ilijar): Switch to using a logger
-            print('Found Detectron ops lib: {}'.format(ops_path))
-            break
-    assert os.path.exists(ops_path), \
-        ('Detectron ops lib not found; make sure that your Caffe2 '
-         'version includes Detectron module')
-    return ops_path
+        for subdir in subdirs:
+            ops_path = os.path.join(prefix, subdir, _DETECTRON_OPS_LIB)
+            if os.path.exists(ops_path):
+                print('Found Detectron ops lib: {}'.format(ops_path))
+                return ops_path
+    raise Exception('Detectron ops lib not found')
 
 
 def get_custom_ops_lib():