diff --git a/docker/RUN_DOCKER_CONTAINER.sh b/docker/RUN_DOCKER_CONTAINER.sh
index 297f83ce..3861daae 100755
--- a/docker/RUN_DOCKER_CONTAINER.sh
+++ b/docker/RUN_DOCKER_CONTAINER.sh
@@ -9,6 +9,11 @@ fi
 
 export WANDB_API_KEY=`cat wandb_key.txt`
 
-docker-compose up ${TAG_NAME} &
-
+if [ ! -z ${DATASET_DIR} ]; then
+  echo "DATASET_DIR is set to ${DATASET_DIR}"
+else
+  export DATASET_DIR=${PWD}/../data
+  echo "DATASET_DIR is not set, using default value instead: ${DATASET_DIR}"
+fi
 
+docker-compose up ${TAG_NAME} &
diff --git a/docker/docker-compose.yml b/docker/docker-compose.yml
index c2258141..3aece008 100644
--- a/docker/docker-compose.yml
+++ b/docker/docker-compose.yml
@@ -15,6 +15,7 @@ services:
     shm_size: '50gb'
     volumes:
       - ${PWD}/../:/root/workspace
+      - ${DATASET_DIR}:/root/workspace/data
       - /tmp/.x11-unix:/tmp/.x11-unix
     environment:
       - DISPLAY=${DISPLAY}