doublelotus commited on
Commit
51c1fe0
·
1 Parent(s): cb92f1e
Files changed (1) hide show
  1. Dockerfile +1 -1
Dockerfile CHANGED
@@ -16,7 +16,7 @@ ENV TRANSFORMERS_CACHE=/code/transformers_cache
16
  ENV HF_HOME=/code/transformers_cache
17
 
18
  # Set PyTorch CUDA memory management variable
19
- ENV PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
20
 
21
  COPY . .
22
 
 
16
  ENV HF_HOME=/code/transformers_cache
17
 
18
  # Set PyTorch CUDA memory management variable
19
+ ENV PYTORCH_CUDA_ALLOC_CONF=garbage_collection_threshold:0.6,max_split_size_mb:128
20
 
21
  COPY . .
22