xref: /aosp_15_r20/external/pytorch/.ci/docker/common/cache_vision_models.sh (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
1#!/bin/bash
2
3set -ex
4
5source "$(dirname "${BASH_SOURCE[0]}")/common_utils.sh"
6
7# Cache the test models at ~/.cache/torch/hub/
8IMPORT_SCRIPT_FILENAME="/tmp/torchvision_import_script.py"
9as_jenkins echo 'import torchvision; torchvision.models.mobilenet_v2(pretrained=True); torchvision.models.mobilenet_v3_large(pretrained=True);' > "${IMPORT_SCRIPT_FILENAME}"
10
11pip_install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cpu
12# Very weird quoting behavior here https://github.com/conda/conda/issues/10972,
13# so echo the command to a file and run the file instead
14conda_run python "${IMPORT_SCRIPT_FILENAME}"
15
16# Cleaning up
17conda_run pip uninstall -y torch torchvision
18rm "${IMPORT_SCRIPT_FILENAME}" || true
19