diff --git a/Dockerfile.gpu b/Dockerfile.gpu
index ba8e04e..eafcfda 100644
--- a/Dockerfile.gpu
+++ b/Dockerfile.gpu
@@ -1,12 +1,11 @@
-FROM nvidia/cuda:10.1-cudnn7-runtime-ubuntu18.04 AS python-ubuntu
+FROM nvidia/cuda:11.2.0-cudnn8-runtime-ubuntu18.04 AS python-ubuntu
 
 RUN apt-get update
 RUN apt-get install -y build-essential curl python3 python3-pip python3-venv
 
 # Inorder to run pip3.6, run it the following way `python3.6 -m pip`
 # Update Python & Install wheel
-RUN python3 -m pip install pip --upgrade
-RUN python3 -m pip install wheel
+RUN python3 -m pip install pip --upgrade && ython3 -m pip install wheel
 
 FROM python-ubuntu
 
@@ -15,7 +14,7 @@ RUN apt-get update &&  apt-get install -y  \
     pkg-config portaudio19-dev swig libpulse-ocaml-dev gfortran \
     libopenblas-dev libatlas-base-dev
 
-RUN mkdir /training && pip3 install precise-lite && pip3 install tensorflow-gpu==2.3.1
+RUN mkdir /training && pip3 install precise-lite && pip3 install tensorflow==2.4.4 && pip3 install tensorflow-gpu==2.4.4
 
 WORKDIR /training
 
diff --git a/entrypoint.sh b/entrypoint.sh
index 555486e..3ba83c5 100755
--- a/entrypoint.sh
+++ b/entrypoint.sh
@@ -2,4 +2,32 @@
 
 set -e
 
-exec $@
\ No newline at end of file
+function populate_env() {
+  if [[ -z "$MODEL_PATH" ] && [ ! -z "$1" ]]; then
+    export MODEL_PATH="$1"
+  fi
+
+  if [[ -z "$DATA_FOLDER" ] && [ ! -z "$2" ]]; then
+    export DATA_FOLDER="$2"
+  fi
+}
+
+populate_env
+
+case "$1" in
+  train)
+    precise-lite-train $TRAIN_ARGS "$MODEL_PATH" "$DATA_FOLDER"
+    ;;
+  "train-incremental")
+    precise-lite-train-incremental $TRAIN_ARGS "$MODEL_PATH" "$DATA_FOLDER"
+    ;;
+  test)
+    precise-lite-test "$MODEL_PATH" "$DATA_FOLDER"
+    ;;
+  convert)
+    precise-lite-convert "$MODEL_PATH"
+    ;;
+  *)
+    exec $@
+    ;;
+esac
\ No newline at end of file