From dc471e88dec851fef65eaba00ec8dfee13ec7ae8 Mon Sep 17 00:00:00 2001
From: bruno <bruno.manga95@gmail.com>
Date: Sat, 14 Nov 2020 12:53:23 +0100
Subject: [PATCH] Add option to build the library using the python's torch
 installation

---
 CMakeLists.txt | 14 ++++++++++++++
 1 file changed, 14 insertions(+)

diff --git a/CMakeLists.txt b/CMakeLists.txt
index 81ca559d530..11f36bacc8b 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -1,9 +1,11 @@
 cmake_minimum_required(VERSION 3.1)
+cmake_policy(SET CMP0074 NEW)
 project(torchvision)
 set(CMAKE_CXX_STANDARD 14)
 set(TORCHVISION_VERSION 0.7.0)
 
 option(WITH_CUDA "Enable CUDA support" OFF)
+option(USE_PYTHON_TORCH "Use python installation of torch" OFF)
 
 if(WITH_CUDA)
   enable_language(CUDA)
@@ -14,6 +16,18 @@ endif()
 
 find_package(Python3 COMPONENTS Development)
 
+if(USE_PYTHON_TORCH)
+  execute_process(COMMAND python3 -c "import torch; print(torch.__file__)"
+    RESULT_VARIABLE COMMAND_RESULT
+    OUTPUT_VARIABLE PYTORCH_INIT_FILE
+    ERROR_VARIABLE PYTHON_ERROR)
+  if(NOT ${COMMAND_RESULT} STREQUAL "0")
+    message(FATAL_ERROR "Failed to find Python's torch module.\n${PYTHON_ERROR}")
+  endif()
+  get_filename_component(Torch_ROOT ${PYTORCH_INIT_FILE} DIRECTORY)
+  message(STATUS "Using Python's torch installation located at ${Torch_ROOT}")
+endif()
+
 find_package(Torch REQUIRED)
 find_package(PNG REQUIRED)
 find_package(JPEG REQUIRED)