From 3597e5e7610ddfac3f5bcbb46ee7a76038ccb0f6 Mon Sep 17 00:00:00 2001
From: James Lamb <jlamb@nvidia.com>
Date: Fri, 18 Oct 2024 08:56:58 -0500
Subject: [PATCH] ensure torch is installed from CUDA index, fix pytest-cov
 restaurant

---
 ci/test_python.sh                     | 2 +-
 ci/test_wheel_pylibwholegraph.sh      | 6 ++++++
 dependencies.yaml                     | 1 -
 python/pylibwholegraph/pyproject.toml | 2 --
 4 files changed, 7 insertions(+), 4 deletions(-)

diff --git a/ci/test_python.sh b/ci/test_python.sh
index a9398c0..3ab5bcc 100755
--- a/ci/test_python.sh
+++ b/ci/test_python.sh
@@ -55,7 +55,7 @@ if [[ "${RUNNER_ARCH}" != "ARM64" ]]; then
   ./ci/run_pylibwholegraph_pytests.sh \
     --junitxml="${RAPIDS_TESTS_DIR}/junit-pylibwholegraph.xml" \
     --cov-config=../../.coveragerc \
-    --cov=cugraph_pyg \
+    --cov=pylibwholegraph \
     --cov-report=xml:"${RAPIDS_COVERAGE_DIR}/pylibwholegraph-coverage.xml" \
     --cov-report=term
 
diff --git a/ci/test_wheel_pylibwholegraph.sh b/ci/test_wheel_pylibwholegraph.sh
index 0898fff..823ae18 100755
--- a/ci/test_wheel_pylibwholegraph.sh
+++ b/ci/test_wheel_pylibwholegraph.sh
@@ -27,6 +27,12 @@ rapids-retry python -m pip install \
     --extra-index-url ${INDEX_URL} \
     "$(echo ./dist/pylibwholegraph*.whl)[test]"
 
+# install torch separately, to be sure we get a CUDA build
+python -m pip install \
+  --index-url "${INDEX_URL}" \
+  -v \
+  'torch>=2.0,<2.4.0a0'
+
 rapids-logger "pytest pylibwholegraph"
 cd python/pylibwholegraph/pylibwholegraph/tests
 python -m pytest \
diff --git a/dependencies.yaml b/dependencies.yaml
index 7c6acc7..c936fe1 100644
--- a/dependencies.yaml
+++ b/dependencies.yaml
@@ -98,7 +98,6 @@ files:
       key: test
     includes:
       - test_python_common
-      - depends_on_pytorch
 
   py_build_cugraph_dgl:
     output: pyproject
diff --git a/python/pylibwholegraph/pyproject.toml b/python/pylibwholegraph/pyproject.toml
index 7e2deec..c2640e5 100644
--- a/python/pylibwholegraph/pyproject.toml
+++ b/python/pylibwholegraph/pyproject.toml
@@ -50,8 +50,6 @@ test = [
     "pytest-forked",
     "pytest-xdist",
     "scipy",
-    "tensordict>=0.1.2",
-    "torch>=2.0,<2.4.0a0",
 ] # This list was generated by `rapids-dependency-file-generator`. To make changes, edit ../../dependencies.yaml and run `rapids-dependency-file-generator`.
 
 [tool.rapids-build-backend]