소스 검색

Upload Dockerfile: expandable_segments fix + reduced params

Sameric 4 달 전
부모
커밋
1bbc0df33e
1개의 변경된 파일2개의 추가작업 그리고 0개의 파일을 삭제
  1. 2 0
      Dockerfile

+ 2 - 0
Dockerfile

@@ -66,6 +66,8 @@ ENV HF_HOME=/home/user/hf_cache
 ENV TRANSFORMERS_CACHE=/home/user/hf_cache
 ENV TORCH_HOME=/home/user/torch_cache
 ENV PATH="/home/user/.local/bin:$PATH"
+# Limit PyTorch CUDA memory fragmentation and allow expandable segments
+ENV PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
 
 USER user