Add sanity checks for torch paths

facebookincubator · Jan 18, 2025 · da8446a · da8446a
1 parent 48491ea
commit da8446a
Show file tree

Hide file tree

Showing 2 changed files with 55 additions and 32 deletions.
diff --git a/pixi.lock b/pixi.lock
diff --git a/pymomentum/CMakeLists.txt b/pymomentum/CMakeLists.txt
@@ -7,13 +7,28 @@
 # Find dependencies
 #===============================================================================
 
+if(NOT DEFINED ENV{CONDA_PREFIX})
+  message(FATAL_ERROR
+    "CONDA_PREFIX is not set. pymomentum currently only supports building"
+    "with Conda/Pixi. Please ensure you are using a Conda or Pixi environment."
+  )
+endif()
+
 set(ENV{NVTOOLSEXT_PATH} "$ENV{CONDA_PREFIX}/include")
+
 if(WIN32)
   set(libtorch_base_path $ENV{CONDA_PREFIX}/Lib/site-packages/torch)
 else()  
   set(libtorch_base_path $ENV{CONDA_PREFIX}/lib/python${Python3_VERSION_MAJOR}.${Python3_VERSION_MINOR}/site-packages/torch)
 endif()
 
+if(NOT EXISTS "${libtorch_base_path}")
+  message(FATAL_ERROR
+    "PyTorch not found in the expected location: ${libtorch_base_path}"
+    "Please ensure PyTorch is installed in your Conda/Pixi environment."
+  )
+endif()
+
 find_package(ATen CONFIG REQUIRED HINTS ${libtorch_base_path})
 find_package(Torch CONFIG REQUIRED HINTS ${libtorch_base_path})
 find_library(torch_python NAMES torch_python HINTS ${libtorch_base_path}/lib/ REQUIRED)