4040 export CUDA_VERSION=" cpu"
4141fi
4242
43- if [[ ${MATRIX_CHANNEL} = ' pypi_release' ]]; then
44- echo " checking pypi release"
45- pip install torch
46- pip install fbgemm-gpu
47- pip install torchrec
48- else
49- # figure out URL
50- if [[ ${MATRIX_CHANNEL} = ' nightly' ]]; then
51- export PYTORCH_URL=" https://download.pytorch.org/whl/nightly/${CUDA_VERSION} "
52- elif [[ ${MATRIX_CHANNEL} = ' test' ]]; then
53- export PYTORCH_URL=" https://download.pytorch.org/whl/test/${CUDA_VERSION} "
54- elif [[ ${MATRIX_CHANNEL} = ' release' ]]; then
55- export PYTORCH_URL=" https://download.pytorch.org/whl/${CUDA_VERSION} "
56- fi
43+ # figure out URL
44+ if [[ ${MATRIX_CHANNEL} = ' nightly' ]]; then
45+ export PYTORCH_URL=" https://download.pytorch.org/whl/nightly/${CUDA_VERSION} "
46+ elif [[ ${MATRIX_CHANNEL} = ' test' ]]; then
47+ export PYTORCH_URL=" https://download.pytorch.org/whl/test/${CUDA_VERSION} "
48+ elif [[ ${MATRIX_CHANNEL} = ' release' ]]; then
49+ export PYTORCH_URL=" https://download.pytorch.org/whl/${CUDA_VERSION} "
50+ fi
5751
58- # install pytorch
59- # switch back to conda once torch nightly is fixed
60- # if [[ ${MATRIX_GPU_ARCH_TYPE} = 'cuda' ]]; then
61- # export PYTORCH_CUDA_PKG="pytorch-cuda=${MATRIX_GPU_ARCH_VERSION}"
62- # fi
63- conda run -n build_binary pip install torch --index-url " $PYTORCH_URL "
52+ # install pytorch
53+ # switch back to conda once torch nightly is fixed
54+ # if [[ ${MATRIX_GPU_ARCH_TYPE} = 'cuda' ]]; then
55+ # export PYTORCH_CUDA_PKG="pytorch-cuda=${MATRIX_GPU_ARCH_VERSION}"
56+ # fi
57+ conda run -n build_binary pip install torch --index-url " $PYTORCH_URL "
6458
65- # install fbgemm
66- conda run -n build_binary pip install fbgemm-gpu --index-url " $PYTORCH_URL "
59+ # install fbgemm
60+ conda run -n build_binary pip install fbgemm-gpu --index-url " $PYTORCH_URL "
6761
68- # install requirements from pypi
69- conda run -n build_binary pip install torchmetrics==1.0.3
62+ # install requirements from pypi
63+ conda run -n build_binary pip install torchmetrics==1.0.3
7064
71- # install torchrec
72- conda run -n build_binary pip install torchrec --index-url " $PYTORCH_URL "
65+ # install torchrec
66+ conda run -n build_binary pip install torchrec --index-url " $PYTORCH_URL "
7367
74- # Run small import test
75- conda run -n build_binary python -c " import torch; import fbgemm_gpu; import torchrec"
76- fi
68+ # Run small import test
69+ conda run -n build_binary python -c " import torch; import fbgemm_gpu; import torchrec"
7770
7871# check directory
7972ls -R
9891
9992if [[ ${MATRIX_CHANNEL} != ' release' ]]; then
10093 exit 0
94+ else
95+ # Check version matches only for release binaries
96+ torchrec_version=$( conda run -n build_binary pip show torchrec | grep Version | cut -d' ' -f2)
97+ fbgemm_version=$( conda run -n build_binary pip show fbgemm_gpu | grep Version | cut -d' ' -f2)
98+
99+ if [ " $torchrec_version " != " $fbgemm_version " ]; then
100+ echo " Error: TorchRec package version does not match FBGEMM package version"
101+ exit 1
102+ fi
101103fi
102104
103105conda create -y -n build_binary python=" ${MATRIX_PYTHON_VERSION} "
104106
105107conda run -n build_binary python --version
106108
107- if [[ ${MATRIX_GPU_ARCH_VERSION} != ' 12.1 ' ]]; then
109+ if [[ ${MATRIX_GPU_ARCH_VERSION} != ' 12.4 ' ]]; then
108110 exit 0
109111fi
110112
@@ -113,6 +115,15 @@ conda run -n build_binary pip install torch
113115conda run -n build_binary pip install fbgemm-gpu
114116conda run -n build_binary pip install torchrec
115117
118+ # Check version matching again for PyPI
119+ torchrec_version=$( conda run -n build_binary pip show torchrec | grep Version | cut -d' ' -f2)
120+ fbgemm_version=$( conda run -n build_binary pip show fbgemm_gpu | grep Version | cut -d' ' -f2)
121+
122+ if [ " $torchrec_version " != " $fbgemm_version " ]; then
123+ echo " Error: TorchRec package version does not match FBGEMM package version"
124+ exit 1
125+ fi
126+
116127# check directory
117128ls -R
118129
0 commit comments