Merge pull request #1807 from shadowdragon89/cp_r2.4_20210204_113128

Cherry pick the potential test fix to 2.4
tensorflow · Feb 4, 2021 · ac3a2eb · ac3a2eb
2 parents 8300bd1 + 644baeb
commit ac3a2eb
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tensorflow_serving/tools/docker/tests/docker_test_lib.sh b/tensorflow_serving/tools/docker/tests/docker_test_lib.sh
@@ -138,7 +138,7 @@ function test_docker_image {
   local docker_opts=" --privileged=true --rm -t -p ${rest_port}:${rest_port}"
   docker_opts+=" -v ${MODEL_FULL_PATH}:${model_base_path}"
   if [ "$USE_NVIDIA_RUNTIME" = true ]; then
-      docker_opts+=" --runtime=nvidia"
+      docker_opts+=" --runtime=nvidia --gpus all"
   fi
 
   if [[ "$IS_MKL_IMAGE" = true ]]; then