NVIDIA · leofang · Apr 29, 2025 · Apr 25, 2025 · Apr 25, 2025 · Apr 25, 2025
diff --git a/cuda_core/tests/example_tests/utils.py b/cuda_core/tests/example_tests/utils.py
@@ -37,7 +37,7 @@ def run_example(samples_path, filename, env=None):
         exec(script, env if env else {})  # nosec B102
     except ImportError as e:
         # for samples requiring any of optional dependencies
-        for m in ("cupy",):
+        for m in ("cupy", "torch"):
             if f"No module named '{m}'" in str(e):
                 pytest.skip(f"{m} not installed, skipping related tests")
                 break

diff --git a/cuda_core/tests/requirements-cu11.txt b/cuda_core/tests/requirements-cu11.txt
@@ -2,3 +2,4 @@ pytest
 # TODO: remove this hack once cupy has a cp313 build
 cupy-cuda11x; python_version < "3.13"
 nvidia-cuda-runtime-cu11  # headers consumed by CuPy
+torch # For PyTorch example
diff --git a/cuda_core/tests/requirements-cu12.txt b/cuda_core/tests/requirements-cu12.txt
@@ -2,3 +2,4 @@ pytest
 # TODO: remove this hack once cupy has a cp313 build
 cupy-cuda12x; python_version < "3.13"
 nvidia-cuda-runtime-cu12  # headers consumed by CuPy
+torch # For PyTorch example