aIbrahiim commented on code in PR #38374:
URL: https://github.com/apache/beam/pull/38374#discussion_r3241795550


##########
.github/workflows/load-tests-pipeline-options/beam_Inference_Python_Benchmarks_Dataflow_Pytorch_Sentiment_Batch_DistilBert_Base_Uncased.txt:
##########
@@ -31,5 +32,7 @@
 --device=CPU
 --input_file=gs://apache-beam-ml/testing/inputs/sentences_50k.txt
 --runner=DataflowRunner
+--sdk_location=container
+--sdk_container_image=us.gcr.io/apache-beam-testing/python-postcommit-it/tensor_rt@sha256:884d67e96d9a3c22fb21fcd412c10a012d4c82a7c723f1c1ffe41fca609b5a6a

Review Comment:
   I needed a CUDA capable worker image so DistilBERT on gpu don't fail with no 
NVIDIA driver / CUDA init on the default dataflow python image so the 
postcommit tensor_rt image is one i already found it used in beam’s GPU IT path 
so i reuse it for the benchmark



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]

Reply via email to