TORCH_USE_CUDA_DSA

Lightning-AI · Mar 21, 2024 · e2b71ac · e2b71ac
1 parent 825cd14
commit e2b71ac
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/dockers/ubuntu-cuda/Dockerfile b/dockers/ubuntu-cuda/Dockerfile
@@ -98,6 +98,7 @@ ENV \
     CUDA_SELECT_NVCC_ARCH_FLAGS="8.0"
 
 ARG TORCH_INSTALL="wheel"
+ENV TORCH_USE_CUDA_DSA=1
 
 RUN \
     if [ "${TORCH_INSTALL}" == "source" ]; then \