Specify correct TF2 container tag

triton-inference-server · Feb 27, 2021 · c7d7b00 · c7d7b00
1 parent 392de0b
commit c7d7b00
Showing 1 changed file with 2 additions and 6 deletions.
diff --git a/build.py b/build.py
@@ -348,12 +348,8 @@ def tensorflow_cmake_args(ver, images):
     if image_name in images:
         image = images[image_name]
     else:
-        if ver == 2:
-            image = 'nvcr.io/nvidia/tensorflow:{}-py3'.format(
-                FLAGS.upstream_container_version)
-        else:
-            image = 'nvcr.io/nvidia/tensorflow:{}-tf{}-py3'.format(
-                FLAGS.upstream_container_version, ver)
+        image = 'nvcr.io/nvidia/tensorflow:{}-tf{}-py3'.format(
+            FLAGS.upstream_container_version, ver)
     return [
         '-DTRITON_TENSORFLOW_VERSION={}'.format(ver),
         '-DTRITON_TENSORFLOW_DOCKER_IMAGE={}'.format(image)