sparkleman commited on
Commit
95dc3fc
·
1 Parent(s): b76f7cc

UPDATE: PYTORCH_CUDA_ALLOC_CONF=expandable_segments

Browse files
Files changed (1) hide show
  1. Dockerfile +2 -1
Dockerfile CHANGED
@@ -50,7 +50,8 @@ ENV HOME=/home/user \
50
  PATH=/usr/local/cuda/bin:/home/user/.local/bin:$PATH \
51
  LD_LIBRARY_PATH="/usr/local/cuda/lib64:${LD_LIBRARY_PATH}" \
52
  CXX=/usr/bin/g++ \
53
- TORCH_CUDA_ARCH_LIST="7.5"
 
54
  WORKDIR $HOME/app
55
 
56
  COPY --chown=user . $HOME/app
 
50
  PATH=/usr/local/cuda/bin:/home/user/.local/bin:$PATH \
51
  LD_LIBRARY_PATH="/usr/local/cuda/lib64:${LD_LIBRARY_PATH}" \
52
  CXX=/usr/bin/g++ \
53
+ TORCH_CUDA_ARCH_LIST="7.5" \
54
+ PYTORCH_CUDA_ALLOC_CONF=expandable_segments
55
  WORKDIR $HOME/app
56
 
57
  COPY --chown=user . $HOME/app