docker note

teddykoker · Mar 24, 2021 · d480043 · d480043
1 parent f7cb735
commit d480043
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/README.md b/README.md
@@ -17,6 +17,11 @@ extension.
 pip install torchsort
 ```
 
+To build the CUDA extension you will need the CUDA toolchain installed. If you
+want to build in an environment without a CUDA runtime (e.g. docker), you will
+need to export the environment variable
+`TORCH_CUDA_ARCH_LIST="Pascal;Volta;Turing"` before installing.
+
 ## Usage
 
 `torchsort` exposes two functions: `soft_rank` and `soft_sort`, each with