NEW: Export CF_TORCH_CUDA_ARCH_LIST to environment

recipe/activate.bat

@@ -0,0 +1,9 @@
+@echo off
+
+if not defined CF_TORCH_CUDA_ARCH_LIST (
+    set "CF_TORCH_CUDA_ARCH_LIST=@cf_torch_cuda_arch_list@"
+    :: "NOT_SET" is used as the value because it is clearer (explicit) that the activation 
+    :: script was run and found no previous value for "CF_TORCH_CUDA_ARCH_LIST".
+    set "CF_TORCH_CUDA_ARCH_LIST_BACKUP=NOT_SET"
+)
+

recipe/activate.sh

@@ -0,0 +1,9 @@
+#!/bin/bash
+
+if [[ ! -v CF_TORCH_CUDA_ARCH_LIST ]]
+then
+    export CF_TORCH_CUDA_ARCH_LIST="@cf_torch_cuda_arch_list@"
+    # "NOT_SET" is used as the value because it is clearer (explicit) that the activation
+    # script was run and found no previous value for "CF_TORCH_CUDA_ARCH_LIST".
+    export CF_TORCH_CUDA_ARCH_LIST_BACKUP="NOT_SET"
+fi

recipe/bld.bat

@@ -108,6 +108,7 @@ if not "%cuda_compiler_version%" == "None" (
     @REM MKLDNN is an Apache-2.0 licensed library for DNNs and is used
     @REM for CPU builds. Not to be confused with MKL.
     set "USE_MKLDNN=1"
+    set "TORCH_CUDA_ARCH_LIST="
 
     @REM On windows, env vars are case-insensitive and setup.py
     @REM passes all env vars starting with CUDA_*, CMAKE_* to
@@ -221,6 +222,21 @@ if "%PKG_NAME%" == "libtorch" (
     @REM Keep the original backed up to sed later
     copy build\CMakeCache.txt build\CMakeCache.txt.orig
     if %ERRORLEVEL% neq 0 exit 1
+
+    if not "%cuda_compiler_version%" == "None" (
+        sed -e "s/@cf_torch_cuda_arch_list@/%TORCH_CUDA_ARCH_LIST%/g" ^
+            %RECIPE_DIR%\activate.bat > %RECIPE_DIR%\activate-replaced.bat
+        if %ERRORLEVEL% neq 0 exit 1
+
+        mkdir %PREFIX%\etc\conda\activate.d
+        copy %RECIPE_DIR%\activate-replaced.bat %PREFIX%\etc\conda\activate.d\libtorch_activate.bat
+        if %ERRORLEVEL% neq 0 exit 1
+
+        mkdir %PREFIX%\etc\conda\deactivate.d
+        copy %RECIPE_DIR%\deactivate.bat %PREFIX%\etc\conda\deactivate.d\libtorch_deactivate.bat
+        if %ERRORLEVEL% neq 0 exit 1
+    )
+
 ) else if "%PKG_NAME%" == "pytorch" (
     @REM Move libtorch_python and remove the other directories afterwards.
     robocopy /NP /NFL /NDL /NJH /E %SP_DIR%\torch\bin\ %LIBRARY_BIN%\ torch_python.dll

recipe/build.sh

@@ -254,6 +254,7 @@ else
     # for CPU builds. Not to be confused with MKL.
     export USE_MKLDNN=1
     export USE_CUDA=0
+    export TORCH_CUDA_ARCH_LIST=""
 fi
 
 echo '${CXX}'=${CXX}
@@ -279,6 +280,16 @@ case ${PKG_NAME} in
 
     # Keep the original backed up to sed later
     cp build/CMakeCache.txt build/CMakeCache.txt.orig
+
+    if [[ "${cuda_compiler_version}" != "None" ]]; then
+        for CHANGE in "activate" "deactivate"
+        do
+            mkdir -p "${PREFIX}/etc/conda/${CHANGE}.d"
+            sed -e "s/@cf_torch_cuda_arch_list@/${TORCH_CUDA_ARCH_LIST}/g" \
+            "${RECIPE_DIR}/${CHANGE}.sh" > "${PREFIX}/etc/conda/${CHANGE}.d/libtorch_${CHANGE}.sh"
+        done
+    fi
+
     ;;
   pytorch)
     $PREFIX/bin/python -m pip install . --no-deps --no-build-isolation -v --no-clean --config-settings=--global-option=-q \

recipe/deactivate.bat

@@ -0,0 +1,6 @@
+@echo off
+
+if "%CF_TORCH_CUDA_ARCH_LIST_BACKUP%" == "NOT_SET" (
+    set "CF_TORCH_CUDA_ARCH_LIST="
+    set "CF_TORCH_CUDA_ARCH_LIST_BACKUP="
+)

recipe/deactivate.sh

@@ -0,0 +1,7 @@
+#!/bin/bash
+
+if [[ "${CF_TORCH_CUDA_ARCH_LIST_BACKUP}" == "NOT_SET" ]]
+then
+  unset CF_TORCH_CUDA_ARCH_LIST
+  unset CF_TORCH_CUDA_ARCH_LIST_BACKUP
+fi