diff --git a/dags/pytorch_xla/configs/pytorchxla_torchbench_config.py b/dags/pytorch_xla/configs/pytorchxla_torchbench_config.py index a229e836..8c333bb2 100644 --- a/dags/pytorch_xla/configs/pytorchxla_torchbench_config.py +++ b/dags/pytorch_xla/configs/pytorchxla_torchbench_config.py @@ -162,8 +162,8 @@ def get_nvidia_driver_install_cmd(driver_version: str) -> str: docker_cmds_ls = ( "apt-get update && apt-get install -y libgl1", - "pip install --user numpy pandas", - "pip install --user --pre torch torchvision torchaudio torchtext -i https://download.pytorch.org/whl/nightly/cu121", + "pip3 install --user numpy pandas", + "pip3 install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu121" "cd /tmp/ && git clone https://github.com/pytorch/benchmark.git", f" cd benchmark && {model_install_cmds()}", "cd /tmp/ && git clone https://github.com/pytorch/pytorch.git",