Enable Python tests for ROCm (facebookresearch#3763)

Summary:
Pull Request resolved: facebookresearch#3763

This change enables Python version of tests for ROCm builds.

Reviewed By: asadoughi

Differential Revision: D61366282

fbshipit-source-id: c2fd688db42d63946f1c5ca7d50f0a1c4d4a33cd

.github/actions/build_cmake/action.yml

@@ -58,7 +58,10 @@ runs:
 
         # install test packages
         conda install -y pytest
-        if [ "${{ inputs.gpu }}" = "ON" ]; then
+        if [ "${{ inputs.rocm }}" = "ON" ]; then
+          : # skip torch install via conda, we need to install via pip to get
+            #  ROCm-enabled version until it's supported in conda by PyTorch
+        elif [ "${{ inputs.gpu }}" = "ON" ]; then
           conda install -y -q pytorch pytorch-cuda=12.4 -c pytorch -c nvidia/label/cuda-12.4.0
         else
           conda install -y -q pytorch -c pytorch
@@ -138,14 +141,19 @@ runs:
       working-directory: build/faiss/python
       run: |
         $CONDA/bin/python setup.py install
+    - name: ROCm - install ROCm-enabled torch via pip
+      if: inputs.rocm == 'ON'
+      shell: bash
+      run: |
+        pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/rocm6.1
     - name: Python tests (CPU only)
       if: inputs.gpu == 'OFF'
       shell: bash
       run: |
         pytest --junitxml=test-results/pytest/results.xml tests/test_*.py
         pytest --junitxml=test-results/pytest/results-torch.xml tests/torch_*.py
     - name: Python tests (CPU + GPU)
-      if: inputs.gpu == 'ON' && inputs.rocm == 'OFF'
+      if: inputs.gpu == 'ON'
       shell: bash
       run: |
         pytest --junitxml=test-results/pytest/results.xml tests/test_*.py
@@ -160,7 +168,6 @@ runs:
         FAISS_DISABLE_CPU_FEATURES=AVX2 LD_DEBUG=libs $CONDA/bin/python -c "import faiss" 2>&1 | grep faiss.so
         LD_DEBUG=libs $CONDA/bin/python -c "import faiss" 2>&1 | grep faiss_avx2.so
     - name: Upload test results
-      if: inputs.rocm == 'OFF'
       uses: actions/upload-artifact@v4
       with:
         name: test-results-arch=${{ runner.arch }}-opt=${{ inputs.opt_level }}-gpu=${{ inputs.gpu }}-raft=${{ inputs.raft }}-rocm=${{ inputs.rocm }}