Update Dockerfile

Dockerfile

@@ -2,7 +2,10 @@ FROM nroggendorff/train-llama:latest
 
 ENV HOME=/home/user \
   HF_HOME=$HOME/.cache/huggingface \
-  PATH=$HOME/.local/bin:$PATH \
-  DATASETS_CACHE=$HOME/.cache/huggingface/datasets \
+  PATH=$HOME/.local/bin:$PATH
+ENV DATASETS_CACHE=$HOME/.cache/huggingface/datasets \
   HF_DATASETS_CACHE=$HOME/.cache/huggingface/datasets \
   TRITON_CACHE_DIR=$HOME/.cache/.triton/
+
+ENV TOKENIZERS_PARALLELISM=false \
+  PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True