Rework the torch-sys build script.

Cargo.toml

@@ -39,7 +39,8 @@ members = ["torch-sys"]
 
 [features]
 default = ["torch-sys/download-libtorch"]
-python = ["cpython"]
+python-libtorch = ["torch-sys/python-libtorch"]
+rl_python = ["cpython"]
 doc-only = ["torch-sys/doc-only"]
 cuda-tests = []
 
@@ -48,7 +49,7 @@ features = [ "doc-only" ]
 
 [[example]]
 name = "reinforcement-learning"
-required-features = ["python"]
+required-features = ["rl_python"]
 
 [[example]]
 name = "stable-diffusion"

README.md

@@ -23,14 +23,22 @@ your system. You can either:
 
 - Use the system-wide libtorch installation (default).
 - Install libtorch manually and let the build script know about it via the `LIBTORCH` environment variable.
-- When a system-wide libtorch can't be found and `LIBTORCH` is not set, the build script will download a pre-built binary version
-of libtorch. By default a CPU version is used. The `TORCH_CUDA_VERSION` environment variable
-can be set to `cu117` in order to get a pre-built binary using CUDA 11.7.
+- Use a Python PyTorch install, to do this set `LIBTORCH_USE_PYTORCH=1`.
+- When a system-wide libtorch can't be found and `LIBTORCH` is not set, the
+  build script will download a pre-built binary version of libtorch. By default
+  a CPU version is used. The `TORCH_CUDA_VERSION` environment variable can be
+  set to `cu117` in order to get a pre-built binary using CUDA 11.7.
 
 ### System-wide Libtorch
 
-The build script will look for a system-wide libtorch library in the following locations:
-- In Linux: `/usr/lib/libtorch.so`
+On linux platforms, the build script will look for a system-wide libtorch
+library in `/usr/lib/libtorch.so`.
+
+### Python PyTorch Install
+
+If the `LIBTORCH_USE_PYTORCH` environment variable is set, the active python
+interpreter is called to retrieve information about the torch python package.
+This version is then linked against.
 
 ### Libtorch Manual Install
 

build.rs

@@ -2,6 +2,9 @@ fn main() {
     let os = std::env::var("CARGO_CFG_TARGET_OS").expect("Unable to get TARGET_OS");
     match os.as_str() {
         "linux" | "windows" => {
+            if let Some(lib_path) = std::env::var_os("DEP_TCH_LIBTORCH_LIB") {
+                println!("cargo:rustc-link-arg=-Wl,-rpath={}", lib_path.to_string_lossy());
+            }
             println!("cargo:rustc-link-arg=-Wl,--no-as-needed");
             println!("cargo:rustc-link-arg=-Wl,--copy-dt-needed-entries");
             println!("cargo:rustc-link-arg=-ltorch");

torch-sys/Cargo.toml

@@ -26,6 +26,7 @@ zip = "0.6"
 [features]
 download-libtorch = ["ureq", "serde", "serde_json"]
 doc-only = []
+python-libtorch = []
 
 [package.metadata.docs.rs]
 features = [ "doc-only" ]