diff --git a/src/build-tf.Dockerfile b/src/build-tf.Dockerfile
index 9052acd..55aaeff 100644
--- a/src/build-tf.Dockerfile
+++ b/src/build-tf.Dockerfile
@@ -1,4 +1,4 @@
-FROM nvaitc/ai-lab:20.01
+FROM nvaitc/ai-lab:20.01-tf2
 
 LABEL maintainer="Timothy Liu <timothyl@nvidia.com>"