Add CUDA environments for pixi (#86)

Summary:
This PR adds multiple environments (by modifying `pixi.toml`) for CPU and GPU environments. Depending on the environment, pymomentum is built with PyTorch with or without CUDA support.

The environment can be set by passing the `-e <env>` option when running pixi tasks. For example, `pixi run -e cpu build_pymomentum` builds without CUDA support while `pixi run -e gpu build_pymomentum` builds with CUDA support. If no `-e` option is passed, the default environment is used, which is `cpu`. To see all available environments, run `pixi info`.

Also, this update includes changes to CMake to find torch packages for any Python versions.

## Checklist:

- [x] Adheres to the [style guidelines](https://facebookincubator.github.io/momentum/docs/developer_guide/style_guide)
- [x] Codebase formatted by running `pixi run lint`

Pull Request resolved: #86

Test Plan:
CI for CPU has been set up. However, the GPU CI setup is not yet complete as I have not figured out how to do it. I am seeking help from the Pixi community.

The GPU environment has been tested locally instead.

Reviewed By: EscapeZero

Differential Revision: D63299584

Pulled By: jeongseok-meta

fbshipit-source-id: 18e583047d7e2d87f3fbb4af34fd835b76557c7c

.github/workflows/ci_ubuntu.yml

@@ -36,15 +36,49 @@ jobs:
           pixi run cmake --build momentum/examples/hello_world/build
 
   pymomentum:
-    name: pymomentum-ubuntu
+    name: pymomentum-${{ matrix.pixi_env }}-ubuntu
     runs-on: ubuntu-latest
+    strategy:
+      fail-fast: false
+      matrix:
+        pixi_env: ["cpu"]  # TODO: Add gpu
     steps:
+      - name: Maximize build space
+        if: ${{ contains(matrix.pixi_env, 'cuda') || contains(matrix.pixi_env, 'gpu') }}
+        uses: easimon/maximize-build-space@master
+        with:
+          root-reserve-mb: 32768
+          remove-dotnet: true
+          remove-android: true
+          remove-haskell: true
+          remove-codeql: true
+          remove-docker-images: true
+
       - name: Checkout
         uses: actions/checkout@v4
+
+      - name: Install CUDA Toolkit
+        if: ${{ contains(matrix.pixi_env, 'cuda') || contains(matrix.pixi_env, 'gpu') }}
+        uses: Jimver/[email protected]
+        id: cuda-toolkit
+        with:
+          cuda: "12.0.1"
+
+      - name: Check CUDA Version
+        if: ${{ contains(matrix.pixi_env, 'cuda') || contains(matrix.pixi_env, 'gpu') }}
+        run: |
+          nvcc --version
+
       - name: Set up pixi
         uses: prefix-dev/[email protected]
         with:
+          environments: ${{ matrix.pixi_env }}
           cache: true
+
+      - name: Check pixi environments
+        run: |
+          pixi info
+
       - name: Build PyMomentum
         run: |
-          pixi run test_pymomentum
+          pixi run -e ${{ matrix.pixi_env }} test_pymomentum