diff --git a/Dockerfile.gpu b/Dockerfile.gpu
index 75325e0..13d763d 100644
--- a/Dockerfile.gpu
+++ b/Dockerfile.gpu
@@ -12,7 +12,7 @@ FROM python-ubuntu
 
 RUN apt-get update &&  apt-get install liblapack-dev libpng-dev libfreetype6-dev libqhull-dev git pkg-config portaudio19-dev swig libpulse-ocaml-dev gfortran libopenblas-dev libatlas-base-dev -y
 
-RUN mkdir /training && pip3 install precise-lite
+RUN mkdir /training && pip3 install precise-lite && pip3 install tensorflow-gpu==2.3.1
 
 WORKDIR /training