Updated environment variables and volumes

docker-compose.yaml

@@ -4,9 +4,12 @@ services:
       context: .
       dockerfile: Dockerfile
     environment:
+      #
       # Set environment variables for your container
-      # PYTORCH_CUDA_ALLOC_CONF: "expandable_segments:True"
-      PYTORCH_CUDA_ALLOC_CONF: max_split_size_mb:128
+      #
+      # expandable_segments (default False) allows the allocator to create a segment initially and then expand its size later when more memory is needed.
+      # Ref: https://pytorch.org/docs/stable/notes/cuda.html
+      PYTORCH_CUDA_ALLOC_CONF: "expandable_segments:False"
     runtime: nvidia  # Use NVIDIA runtime for GPU support
     deploy:
       resources: