diff --git a/ubuntuWithNvidiaDriver/.gitignore b/ubuntuWithNvidiaDriver/.gitignore
new file mode 100644
index 0000000..522719e
--- /dev/null
+++ b/ubuntuWithNvidiaDriver/.gitignore
@@ -0,0 +1 @@
+NVIDIA-DRIVER.run
\ No newline at end of file
diff --git a/ubuntuWithNvidiaDriver/Dockerfile b/ubuntuWithNvidiaDriver/Dockerfile
index 1790e4f..5f7d498 100644
--- a/ubuntuWithNvidiaDriver/Dockerfile
+++ b/ubuntuWithNvidiaDriver/Dockerfile
@@ -3,9 +3,12 @@ MAINTAINER github/gklingler
 
 RUN apt-get update
 RUN apt-get install -y mesa-utils
+RUN apt-get install -y module-init-tools
+RUN apt-get install -y x11-utils
 
-# install nvidia driver
+# Install nvidia driver
 RUN apt-get install -y binutils
 ADD NVIDIA-DRIVER.run /tmp/NVIDIA-DRIVER.run
-RUN sh /tmp/NVIDIA-DRIVER.run -a -N --ui=none --no-kernel-module
+# Docker kernel already contains nvidia modules, all we need are the corresponding files
+RUN sh /tmp/NVIDIA-DRIVER.run -a --ui=none --no-kernel-module
 RUN rm /tmp/NVIDIA-DRIVER.run
diff --git a/ubuntuWithNvidiaDriver/build.sh b/ubuntuWithNvidiaDriver/build.sh
index b80b939..3953ce3 100755
--- a/ubuntuWithNvidiaDriver/build.sh
+++ b/ubuntuWithNvidiaDriver/build.sh
@@ -1,7 +1,19 @@
 #!/bin/sh
 
 IMAGE_NAME=ubuntu_with_nvidia_driver
-NVIDIA_DRIVER=NVIDIA-Linux-x86_64-337.12.run  # path to nvidia driver
+NVIDIA_DRIVER=NVIDIA-Linux-x86_64-*.run  # path to nvidia driver
 
-cp ${NVIDIA_DRIVER} NVIDIA-DRIVER.run
-sudo docker build -t ${IMAGE_NAME} .
+nvidia_host_version() {
+    echo "$(head -n1 </proc/driver/nvidia/version | cut -d' ' -f 9)"
+}
+
+linux_arch() {
+    echo "$(uname)-$(arch)"
+}
+
+nvidia_url() {
+    echo "https://http.download.nvidia.com/XFree86/$(linux_arch)/$(nvidia_host_version)/NVIDIA-$(linux_arch)-$(nvidia_host_version).run"
+}
+
+[ -f "NVIDIA-DRIVER.run" ] || wget $(nvidia_url) -O NVIDIA-DRIVER.run || echo "Please download nvidia driver with same version as the host. Probably $(nvidia_host_version). Guessed URL $(nvidia_url) but it did not work"
+docker build -t ${IMAGE_NAME} .
diff --git a/ubuntuWithNvidiaDriver/run.sh b/ubuntuWithNvidiaDriver/run.sh
index 9fec378..636ba4b 100755
--- a/ubuntuWithNvidiaDriver/run.sh
+++ b/ubuntuWithNvidiaDriver/run.sh
@@ -2,5 +2,12 @@
 
 COMMAND=/bin/bash
 
-xhost + # allow connections to X server
-docker run --privileged -e "DISPLAY=unix:0.0" -v="/tmp/.X11-unix:/tmp/.X11-unix:rw"  -i -t ubuntu_with_nvidia_driver $COMMAND
+# allows docker to connect to x11 server at the host
+XAUTH=/tmp/.docker.xauth
+touch $XAUTH
+xauth nlist $DISPLAY | sed -e 's/^..../ffff/' | xauth -f $XAUTH nmerge -
+
+
+docker run --privileged -e "DISPLAY=unix:0.0" -v="/tmp/.X11-unix:/tmp/.X11-unix:rw"  \
+       -e "XAUTHORITY=$XAUTH" -v="$XAUTH:$XAUTH:rw" \
+       -it ubuntu_with_nvidia_driver $COMMAND