[Electra/TF2] Fix container image and num_gpu arg (NVIDIA#628)

* Fix container image * fix num_gpu arg
fumikawa · Jul 30, 2020 · 5a87078 · 5a87078
1 parent a275790
commit 5a87078
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 3 deletions.
diff --git a/TensorFlow2/LanguageModeling/ELECTRA/Dockerfile b/TensorFlow2/LanguageModeling/ELECTRA/Dockerfile
@@ -11,8 +11,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-ARG FROM_IMAGE_NAME=nvcr.io/nvidian/tensorflow:20.06-tf2-py3
-#FROM gitlab-master.nvidia.com:5005/dl/dgx/tensorrtserver:master-py3.1164446-client as trt
+ARG FROM_IMAGE_NAME=nvcr.io/nvidia/tensorflow:20.06-tf2-py3
 #FROM nvcr.io/nvidia/tensorrtserver:20.03-py3-clientsdk as trt
 FROM ${FROM_IMAGE_NAME}
 RUN apt-get update && apt-get install -y pbzip2 pv bzip2 cabextract

diff --git a/TensorFlow2/LanguageModeling/ELECTRA/scripts/benchmark.sh b/TensorFlow2/LanguageModeling/ELECTRA/scripts/benchmark.sh
@@ -25,4 +25,4 @@ OUT_DIR=${8:-"results/"}
 init_checkpoint=${9:-"None"}
 cache_dir=${10:-"$squad_dir"}
 
-bash scripts/run_squad.sh google/electra-base-discriminator 1 $batch_size $infer_batch_size 8e-4 $precision 8 $RANDOM $SQUAD_VERSION $squad_dir $OUT_DIR $init_checkpoint $mode interactive $cache_dir 200
+bash scripts/run_squad.sh google/electra-base-discriminator 1 $batch_size $infer_batch_size 8e-4 $precision $num_gpu $RANDOM $SQUAD_VERSION $squad_dir $OUT_DIR $init_checkpoint $mode interactive $cache_dir 200