trying to avoid out of memory issues (#6)

valeriano-manassero · Mar 10, 2020 · c300629 · c300629
1 parent 4242f3e
commit c300629
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/Dockerfile b/Dockerfile
@@ -9,6 +9,7 @@ ENV NVIDIA_DRIVER_CAPABILITIES compute,utility
 ENV NVIDIA_REQUIRE_CUDA "cuda>=10.1 brand=tesla,driver>=384,driver<385 brand=tesla,driver>=396,driver<397 brand=tesla,driver>=410,driver<411"
 ENV NCCL_VERSION 2.4.8
 ENV CUDNN_VERSION 7.6.5.32
+ENV TF_FORCE_GPU_ALLOW_GROWTH true
 
 RUN apt-get update && apt-get install -y --no-install-recommends ca-certificates apt-transport-https gnupg-curl && \
     NVIDIA_GPGKEY_SUM=d1be581509378368edeec8c1eb2958702feedf3bc3d17011adbf24efacce4ab5 && \