python3 only when available

Signed-off-by: Eli Uriegas <eliuriegas@fb.com>
pytorch · Jul 26, 2022 · a30b258 · a30b258
1 parent 6e2edb1
commit a30b258
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/packaging/pkg_helpers.bash b/packaging/pkg_helpers.bash
@@ -221,8 +221,13 @@ setup_pip_pytorch_version() {
 setup_conda_pytorch_constraint() {
   if [[ -z "$PYTORCH_VERSION" ]]; then
     export CONDA_CHANNEL_FLAGS="${CONDA_CHANNEL_FLAGS} -c pytorch-nightly -c pytorch"
+    PYTHON="python"
+    # Check if we have python 3 instead and prefer that
+    if python3 --version >/dev/null 2>/dev/null; then
+      PYTHON="python3"
+    fi
     export PYTORCH_VERSION="$(conda search --json 'pytorch[channel=pytorch-nightly]' | \
-                              python3 -c "import os, sys, json, re; cuver = os.environ.get('CU_VERSION'); \
+                              ${PYTHON} -c "import os, sys, json, re; cuver = os.environ.get('CU_VERSION'); \
                                cuver_1 = cuver.replace('cu', 'cuda') if cuver != 'cpu' else cuver; \
                                cuver_2 = (cuver[:-1] + '.' + cuver[-1]).replace('cu', 'cuda') if cuver != 'cpu' else cuver; \
                                print(re.sub(r'\\+.*$', '', \