xref: /aosp_15_r20/external/executorch/.ci/docker/common/install_pytorch.sh (revision 523fa7a60841cd1ecfb9cc4201f1ca8b03ed023a)
1#!/bin/bash
2# Copyright (c) Meta Platforms, Inc. and affiliates.
3# All rights reserved.
4#
5# This source code is licensed under the BSD-style license found in the
6# LICENSE file in the root directory of this source tree.
7
8set -ex
9
10# shellcheck source=/dev/null
11source "$(dirname "${BASH_SOURCE[0]}")/utils.sh"
12
13install_domains() {
14  echo "Install torchvision and torchaudio"
15  pip_install --no-use-pep517 --user "git+https://github.com/pytorch/audio.git@${TORCHAUDIO_VERSION}"
16  pip_install --no-use-pep517 --user "git+https://github.com/pytorch/vision.git@${TORCHVISION_VERSION}"
17}
18
19install_pytorch_and_domains() {
20  git clone https://github.com/pytorch/pytorch.git
21
22  # Fetch the target commit
23  pushd pytorch || true
24  git checkout "${TORCH_VERSION}"
25  git submodule update --init --recursive
26
27  chown -R ci-user .
28
29  export _GLIBCXX_USE_CXX11_ABI=0
30  # Then build and install PyTorch
31  conda_run python setup.py bdist_wheel
32  pip_install "$(echo dist/*.whl)"
33
34  # Grab the pinned audio and vision commits from PyTorch
35  TORCHAUDIO_VERSION=$(cat .github/ci_commit_pins/audio.txt)
36  export TORCHAUDIO_VERSION
37  TORCHVISION_VERSION=$(cat .github/ci_commit_pins/vision.txt)
38  export TORCHVISION_VERSION
39
40  install_domains
41
42  popd || true
43  # Clean up the cloned PyTorch repo to reduce the Docker image size
44  rm -rf pytorch
45
46  # Print sccache stats for debugging
47  as_ci_user sccache --show-stats
48}
49
50install_pytorch_and_domains
51