diff --git a/.gitignore b/.gitignore
index 5c82ba8..e3e0c0c 100644
--- a/.gitignore
+++ b/.gitignore
@@ -20,3 +20,4 @@ __pycache__
 ._*
 
 slurm-*.out
+.idea/*
diff --git a/Dockerfile.gpu b/Dockerfile.gpu
index 179db5c..6dcaf71 100644
--- a/Dockerfile.gpu
+++ b/Dockerfile.gpu
@@ -1,53 +1,32 @@
-FROM nvidia/cuda:9.0-base-ubuntu16.04
-ARG DEBIAN_FRONTEND="noninteractive"
+FROM tensorflow/tensorflow:1.12.3-gpu-py3
 ENV LANG="C.UTF-8" \
     LC_ALL="C.UTF-8"
 
-# Install GPU dependencies for TensorFlow.
+# remove old nvidia keys and install new signing keys
+RUN rm /etc/apt/sources.list.d/cuda.list
+RUN rm /etc/apt/sources.list.d/nvidia-ml.list
 RUN apt-get update \
-    && apt-get install -y --no-install-recommends \
-        nvinfer-runtime-trt-repo-ubuntu1604-4.0.1-ga-cuda9.0 \
-    && apt-get update \
-    && apt-get install -y --no-install-recommends \
-        libnvinfer4=4.1.2-1+cuda9.0 \
-        cuda-cublas-9-0 \
-        cuda-cufft-9-0 \
-        cuda-curand-9-0 \
-        cuda-cusolver-9-0 \
-        cuda-cusparse-9-0 \
-        libcudnn7=7.2.1.38-1+cuda9.0 \
-        libnccl2=2.2.13-1+cuda9.0 \
-        libfreetype6-dev \
-        libhdf5-serial-dev \
-        libpng12-dev \
-        libzmq3-dev \
-        pkg-config \
-        unzip \
-    && apt-get clean \
-    && rm -rf /var/lib/apt/lists/*
+    && apt-key del 7fa2af80 \
+    && apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu1604/x86_64/3bf863cc.pub \
+    && apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/machine-learning/repos/ubuntu1604/x86_64/7fa2af80.pub
 
 RUN apt-get update \
     && apt-get install --yes --quiet --no-install-recommends \
-        ca-certificates \
-        curl \
         git \
-        libgomp1 \
-        python3 \
+        wget \
     && rm -rf /var/lib/apt/lists/* \
-    && curl -fsSL https://bootstrap.pypa.io/pip/3.5/get-pip.py | python3 - \
-    && ln -s $(which python3) /usr/local/bin/python
+    && curl -fsSL https://bootstrap.pypa.io/pip/3.5/get-pip.py | python3 -
 
 WORKDIR /opt/kwyk
 COPY [".", "."]
-RUN pip install --no-cache-dir --editable .[gpu] \
+RUN python3 -m pip install --no-cache-dir --editable .[gpu] \
     # Install the saved models.
     && curl -fsSL https://github.com/patrick-mcclure/nobrainer/tarball/master \
     | tar xz --strip=1 --wildcards '*/saved_models'
 
-
 ENV FREESURFER_HOME="/opt/kwyk/freesurfer"
 ENV PATH="$FREESURFER_HOME/bin:$PATH"
 
 WORKDIR /data
 ENTRYPOINT ["kwyk"]
-LABEL maintainer="Jakub Kaczmarzyk <jakub.kaczmarzyk@gmail.com>"
+
diff --git a/setup.cfg b/setup.cfg
index fc6d384..1b49024 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -30,8 +30,8 @@ console_scripts =
     kwyk = kwyk.cli:predict
 
 [options.extras_require]
-cpu = tensorflow == 1.12.3
-gpu = tensorflow-gpu == 1.12.3
+cpu = tensorflow >= 1.12.0
+gpu = tensorflow-gpu >= 1.12.0
 
 [versioneer]
 VCS = git