⬆️ Update for TensorFlow 2.5+2.6

assets/Dockerfile

@@ -2,7 +2,9 @@ ARG BASE_IMAGE
 
 FROM $BASE_IMAGE
 
-ARG tf
+ARG tensorflow_package
+ARG tensorflow_version
+ARG test_nvidia_options
 
 RUN apt-get update && \
     apt-get install -y python3 python3-distutils curl &&\
@@ -11,9 +13,9 @@ RUN apt-get update && \
     curl -sSL https://bootstrap.pypa.io/get-pip.py -o get-pip.py && \
     python3 get-pip.py && \
     rm -f get-pip.py
-COPY assets/requirements-$tf.txt /tmp
-RUN pip install --no-cache-dir -r /tmp/requirements-$tf.txt
+RUN pip install --no-cache-dir "$tensorflow_package == $tensorflow_version"
 
 COPY test-nvidia /usr/bin
 
-CMD ["/usr/bin/test-nvidia"]
+ENV test_nvidia_options $test_nvidia_options
+CMD /usr/bin/test-nvidia $test_nvidia_options

assets/requirements-tf1.txt

@@ -1 +0,0 @@
-tensorflow-gpu == 1.15.*

assets/requirements-tf2.txt

@@ -1 +0,0 @@
-tensorflow == 2.*