Add cuda-cudart-dev and cuda-profiler-api to environment (#549)

Otherwise, I suddenly cannot build apex (for running DLRM) anymore...
eth-easl · Jun 25, 2024 · 70ce04c · 70ce04c
1 parent b1b71c2
commit 70ce04c
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 0 deletions.
diff --git a/environment.yml b/environment.yml
@@ -47,3 +47,5 @@ dependencies:
 #  - nvidia::cuda-nvcc=12.1.*
 #  - nvidia::cuda-nvtx=12.1.*
 #  - nvidia::cuda-cupti=12.1.*
+#  - nvidia::cuda-cudart-dev=12.1.*
+#  - nvidia::cuda-profiler-api=12.1.*
diff --git a/scripts/initial_setup.sh b/scripts/initial_setup.sh
@@ -76,6 +76,8 @@ if [ "$IS_MAC" != true ]; then
         sed -i "/nvidia::cuda-nvcc/c\  - nvidia::cuda-nvcc=$CUDA_VERSION.*" $PARENT_DIR/environment.yml
         sed -i "/nvidia::cuda-nvtx/c\  - nvidia::cuda-nvtx=$CUDA_VERSION.*" $PARENT_DIR/environment.yml
         sed -i "/nvidia::cuda-cupti/c\  - nvidia::cuda-cupti=$CUDA_VERSION.*" $PARENT_DIR/environment.yml
+        sed -i "/nvidia::cuda-cudart-dev/c\  - nvidia::cuda-cudart-dev=$CUDA_VERSION.*" $PARENT_DIR/environment.yml
+        sed -i "/nvidia::cuda-profiler-api=/c\  - nvidia::cuda-profiler-api==$CUDA_VERSION.*" $PARENT_DIR/environment.yml
 
         # Fix $PARENT_DIR/docker-compose.yml
         startLine="$(grep -n "CUDASTART" $PARENT_DIR/docker-compose.yml | head -n 1 | cut -d: -f1)"