diff --git a/cuda9.1-python3.6-tensorflow/Dockerfile b/cuda9.1-python3.6-tensorflow/Dockerfile
index 0dbd11f..84c620a 100644
--- a/cuda9.1-python3.6-tensorflow/Dockerfile
+++ b/cuda9.1-python3.6-tensorflow/Dockerfile
@@ -1,6 +1,4 @@
-# FROM tiangolo/python-machine-learning:cuda9.1-python3.6-tensorflow
-# TODO fix this
-FROM tiangolo/python-machine-learning:cuda9.1-python3.6
+FROM tiangolo/python-machine-learning:cuda9.1-python3.6-tensorflow
 
 LABEL maintainer="Sebastian Ramirez <tiangolo@gmail.com>"