enable unit test gpu (#2228)

Summary:
Pull Request resolved: #2228

# context

Differential Revision: D51095381

.github/workflows/unittest_ci.yml

@@ -5,18 +5,18 @@ name: Unit Test CI
 
 on:
   # TODO: re-enable when GPU unit tests are working
-  # push:
-  #   paths-ignore:
-  #     - "docs/*"
-  #     - "third_party/*"
-  #     - .gitignore
-  #     - "*.md"
-  # pull_request:
-  #   paths-ignore:
-  #     - "docs/*"
-  #     - "third_party/*"
-  #     - .gitignore
-  #     - "*.md"
+  push:
+    paths-ignore:
+      - "docs/*"
+      - "third_party/*"
+      - .gitignore
+      - "*.md"
+  pull_request:
+    paths-ignore:
+      - "docs/*"
+      - "third_party/*"
+      - .gitignore
+      - "*.md"
   workflow_dispatch:
 
 jobs:
@@ -29,19 +29,19 @@ jobs:
          - os: linux.2xlarge
           # ideally we run on 3.9 and 3.10 as well, however we are limited in resources.
            python-version: 3.8
-           python-tag: "py38"
-           cuda-tag: "cu11"
+           python-tag: "py39"
+           cuda-tag: "cu12"
     steps:
     # Checkout the repository to the GitHub Actions runner
     - name: Check ldd --version
       run: ldd --version
     - name: Checkout
-      uses: actions/checkout@v2
+      uses: actions/checkout@v3
     - name: Update pip
       run: |
         sudo yum update -y
         sudo yum -y install git python3-pip
-        sudo pip3 install --upgrade pip
+        # sudo pip3 install --upgrade pip
     - name: Setup conda
       run: |
         wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh -O ~/miniconda.sh
@@ -99,7 +99,7 @@ jobs:
       matrix:
         os: [linux.4xlarge.nvidia.gpu]
         python-version: [3.8]
-        cuda-tag: ["cu11"]
+        cuda-tag: ["cu12"]
     needs: build_on_cpu
     # the glibc version should match the version of the one we used to build the binary
     # for this case, it's 2.26
@@ -135,7 +135,7 @@ jobs:
         sudo lshw -C display
     # Checkout the repository to the GitHub Actions runner
     - name: Checkout
-      uses: actions/checkout@v2
+      uses: actions/checkout@v3
     - name: Update pip
       run: |
         sudo yum update -y
@@ -170,7 +170,7 @@ jobs:
       shell: bash
       run: |
         conda run -n build_binary \
-          pip install fbgemm-gpu --index-url https://download.pytorch.org/whl/nightly/cu118
+          pip install fbgemm-gpu --index-url https://download.pytorch.org/whl/nightly/cu12
     # download wheel from GHA
     - name: Download wheel
       uses: actions/download-artifact@v2
@@ -181,7 +181,7 @@ jobs:
     - name: Install TorchRec GPU
       run: |
         rm -r dist || true
-        conda run -n build_binary python -m pip install dist/*.whl
+        conda run -n build_binary python -m pip install *.whl
     - name: Test torchrec installation
       shell: bash
       run: |