chore: documenation building fixes

Author: narendasan

.github/workflows/docgen.yml

@@ -13,15 +13,17 @@ jobs:
     runs-on: linux.g5.4xlarge.nvidia.gpu
     container:
       image: docker.io/pytorch/manylinux-builder:cuda12.1
+    env:
+      CUDA_HOME: /usr/local/cuda-12.1
+      VERSION_SUFFIX: cu121
+      CI_BUILD: 1
     steps:
       - uses: actions/checkout@v3
         with:
           ref: ${{github.head_ref}}
       - name: Select Python / CUDA
         run: |
           echo "/opt/python/cp311-cp311/bin/" >> $GITHUB_PATH
-          echo "CUDA_HOME=/usr/local/cuda-12.1" >> $GITHUB_ENV
-          echo "VERSION_SUFFIX=cu121" >> $GITHUB_ENV
       - name: Install base deps
         run: |
           python3 -m pip install pyyaml torch --pre --extra-index-url https://download.pytorch.org/whl/nightly/cu121
@@ -32,7 +34,7 @@ jobs:
       - name: Build Python Package
         run: |
           python3 -m pip install pip --upgrade
-          CI_BUILD=1 python3 -m pip install --pre . --extra-index-url https://download.pytorch.org/whl/nightly/cu121
+          python3 -m pip install --pre . --extra-index-url https://download.pytorch.org/whl/nightly/cu121
       - name: Generate New Docs
         run: |
           cd docsrc