瀏覽代碼

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