fix: link PyTorch libraries in wheel builds

.github/workflows/build.yml

@@ -8,19 +8,10 @@ on:
   workflow_dispatch:
 
 jobs:
-  # Build for Ubuntu 22.04 + Python 3.11 with explicit CPU/CUDA build-path matrix
+  # Build the distributable CPU wheel on a standard hosted runner.
   build:
-    name: Build (${{ matrix.build_mode }}, Ubuntu 22.04, Python 3.11)
+    name: Build (CPU, Ubuntu 22.04, Python 3.11)
     runs-on: ubuntu-22.04
-    strategy:
-      fail-fast: false
-      matrix:
-        build_mode: [cpu, cuda]
-        include:
-          - build_mode: cpu
-            amms_enable_cuda: "0"
-          - build_mode: cuda
-            amms_enable_cuda: "1"
 
     steps:
       - name: Checkout code
@@ -39,26 +30,16 @@ jobs:
       - name: Install Python dependencies
         run: |
           python -m pip install --upgrade pip
-          pip install build wheel setuptools>=64
-          pip install torch>=2.0.0 numpy>=1.20.0
+          pip install build wheel "setuptools>=64" "cmake>=3.14" "pybind11>=2.10.0"
+          pip install torch==2.5.1 --index-url https://download.pytorch.org/whl/cpu
+          pip install "numpy>=1.20.0"
           pip install pytest pytest-cov
 
-      - name: Prepare CUDA toolchain switch path
-        if: matrix.build_mode == 'cuda'
-        run: |
-          mkdir -p /tmp/fake-cuda/bin
-          cat > /tmp/fake-cuda/bin/nvcc << 'EOF'
-          #!/usr/bin/env bash
-          echo "fake nvcc for AMMS_ENABLE_CUDA switch validation"
-          EOF
-          chmod +x /tmp/fake-cuda/bin/nvcc
-
       - name: Build package
         run: |
-          python -m build --wheel
+          python -m build --wheel --no-isolation
         env:
-          AMMS_ENABLE_CUDA: ${{ matrix.amms_enable_cuda }}
-          CUDA_HOME: /tmp/fake-cuda
+          AMMS_ENABLE_CUDA: "0"
 
       - name: Install package
         run: |
@@ -78,9 +59,28 @@ jobs:
       - name: Upload wheel
         uses: actions/upload-artifact@v4
         with:
-          name: wheel-${{ matrix.build_mode }}-ubuntu22.04-py311
+          name: wheel-cpu-ubuntu22.04-py311
           path: dist/*.whl
 
+  cuda-contract:
+    name: CUDA build contract
+    runs-on: ubuntu-22.04
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v4
+
+      - name: Set up Python 3.11
+        uses: actions/setup-python@v5
+        with:
+          python-version: "3.11"
+
+      - name: Validate explicit CUDA switch
+        run: |
+          pip install pytest numpy
+          pytest \
+            tests/test_issue5_cuda_cpu_switch_cleanup.py \
+            tests/test_issue6_build_matrix_and_perf_baseline.py::test_issue6_setup_has_explicit_cuda_cpu_switch_contract
+
   # Code quality checks
   lint:
     name: Code Quality

sage/libs/amms/implementations/CMakeLists.txt

@@ -17,8 +17,9 @@ include(cmake/default.cmake)
 include(cmake/FindCuda.cmake)
 include(cmake/FindTorch.cmake)
 
-# Find PyTorch
-find_package(Torch REQUIRED)
+# Find PyTorch through the package's official CMake config. Using CONFIG avoids
+# shadowing it with the local hint-only cmake/FindTorch.cmake module.
+find_package(Torch REQUIRED CONFIG)
 message(STATUS "Found Torch version: ${Torch_VERSION}")
 message(STATUS "Torch libraries: ${TORCH_LIBRARIES}")
 message(STATUS "Torch include directories: ${TORCH_INCLUDE_DIRS}")
@@ -279,4 +280,4 @@ if(NOT LIBAMM_SKIP_DATASET_COPY AND EXISTS "${LIBAMM_DATASET_SOURCE_DIR}")
     file(COPY "${LIBAMM_DATASET_SOURCE_DIR}/" DESTINATION "${CMAKE_BINARY_DIR}/benchmark/datasets/")
 else()
     message(STATUS "Dataset copy: SKIPPED")
-endif()
+endif()

sage/libs/amms/implementations/src/PyAMM.cpp

@@ -9,6 +9,7 @@
 #include <pybind11/stl.h>
 #include <pybind11/numpy.h>
 #include <torch/torch.h>  // Still needed internally for LibAMM
+#include <cstring>
 #include <Utils/ConfigMap.hpp>
 #include <Utils/IntelliLog.h>
 #include <LibAMM.h>
@@ -33,14 +34,6 @@ py::array torch_to_numpy(const torch::Tensor& tensor) {
 
     // Get tensor properties
     auto sizes = cpu_tensor.sizes();
-    auto strides = cpu_tensor.strides();
-
-    // Convert strides from elements to bytes
-    std::vector<ssize_t> np_strides;
-    for (auto s : strides) {
-        np_strides.push_back(s * cpu_tensor.element_size());
-    }
-
     std::vector<ssize_t> np_shape(sizes.begin(), sizes.end());
 
     // Determine NumPy dtype based on torch dtype
@@ -57,9 +50,11 @@ py::array torch_to_numpy(const torch::Tensor& tensor) {
         throw std::runtime_error("Unsupported tensor dtype for NumPy conversion");
     }
 
-    // Create NumPy array sharing the same memory (zero-copy when possible)
-    // Note: We need to keep the tensor alive, so we make a copy for safety
-    return py::array(dtype, np_shape, np_strides, cpu_tensor.data_ptr(), py::cast(cpu_tensor));
+    // Copy into NumPy-owned memory so the result does not depend on a registered
+    // pybind caster or on the lifetime of the temporary tensor.
+    py::array result(dtype, np_shape);
+    std::memcpy(result.mutable_data(), cpu_tensor.data_ptr(), cpu_tensor.nbytes());
+    return result;
 }
 
 // Convert NumPy array to torch::Tensor (LibAMM-facing)

tests/test_issue4_boundary_cleanup.py

@@ -2,7 +2,6 @@
 
 from pathlib import Path
 
-
 REPO_ROOT = Path(__file__).parent.parent
 
 

tests/test_issue6_build_matrix_and_perf_baseline.py

@@ -10,17 +10,17 @@
 import pytest
 
 
-def test_issue6_build_workflow_has_cpu_cuda_matrix() -> None:
-    """CI workflow includes explicit CPU/CUDA build-path matrix coverage."""
+def test_issue6_build_workflow_has_cpu_build_and_cuda_contract() -> None:
+    """CI builds a CPU wheel and validates the explicit CUDA switch contract."""
     repo_root = Path(__file__).parent.parent
     workflow_content = (repo_root / ".github" / "workflows" / "build.yml").read_text(
         encoding="utf-8"
     )
 
-    assert "build_mode: [cpu, cuda]" in workflow_content
-    assert "amms_enable_cuda" in workflow_content
-    assert "AMMS_ENABLE_CUDA: ${{ matrix.amms_enable_cuda }}" in workflow_content
-    assert "Prepare CUDA toolchain switch path" in workflow_content
+    assert "Build (CPU, Ubuntu 22.04, Python 3.11)" in workflow_content
+    assert 'AMMS_ENABLE_CUDA: "0"' in workflow_content
+    assert "cuda-contract:" in workflow_content
+    assert "CUDA build contract" in workflow_content
 
 
 def test_issue6_setup_has_explicit_cuda_cpu_switch_contract() -> None: