enable pytorch ops in cuda test

deepmodeling · Jun 20, 2024 · ee0f409 · ee0f409
1 parent 83d71ac
commit ee0f409
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/.github/workflows/test_cuda.yml b/.github/workflows/test_cuda.yml
@@ -49,11 +49,13 @@ jobs:
     - run: python -m pip install -U uv
     - run: python -m uv pip install --system "tensorflow>=2.15.0rc0" "torch>=2.2.0"
     - run: |
+        export PYTORCH_ROOT=$(python -c 'import torch;print(torch.__path__[0])')
         export TENSORFLOW_ROOT=$(python -c 'import importlib,pathlib;print(pathlib.Path(importlib.util.find_spec("tensorflow").origin).parent)')
         python -m uv pip install --system -v -e .[gpu,test,lmp,cu12,torch] mpi4py
       env:
         DP_VARIANT: cuda
         DP_ENABLE_NATIVE_OPTIMIZATION: 1
+        DP_ENABLE_PYTORCH: 1
     - run: dp --version
     - run: python -m pytest source/tests --durations=0
       env: