diff --git a/example/pod/alexnet-gpu.yaml b/example/pod/alexnet-gpu.yaml index 37fb96b28fbe04031ff6579c6af0af08111e99a7..f1c2519c2e551100acf918a2fb7b85b0bf408988 100644 --- a/example/pod/alexnet-gpu.yaml +++ b/example/pod/alexnet-gpu.yaml @@ -7,13 +7,13 @@ metadata: spec: containers: - name: alexnet-tf-gpu-container - image: rocm/tensorflow:rocm1.7.1 + image: rocm/tensorflow:latest workingDir: /root env: - name: HIP_VISIBLE_DEVICES value: "0" # # 0,1,2,...,n for running on GPU and select the GPUs, -1 for running on CPU command: ["/bin/bash", "-c", "--"] - args: ["python convnet-benchmarks/tensorflow/benchmark_alexnet.py; trap : TERM INT; sleep infinity & wait"] + args: ["python3 benchmarks/scripts/tf_cnn_benchmarks/tf_cnn_benchmarks.py --model=alexnet; trap : TERM INT; sleep infinity & wait"] resources: limits: amd.com/gpu: 1 # requesting a GPU