From 12b88204b000547baf87e8575bcc94eb5ec2fef2 Mon Sep 17 00:00:00 2001
From: Nicolas Patry <patry.nicolas@protonmail.com>
Date: Sat, 14 Sep 2024 18:43:12 +0200
Subject: [PATCH] Putting the cuda package in the flake.

---
 .github/workflows/nix_integration_tests.yaml | 5 -----
 flake.nix                                    | 1 +
 2 files changed, 1 insertion(+), 5 deletions(-)

diff --git a/.github/workflows/nix_integration_tests.yaml b/.github/workflows/nix_integration_tests.yaml
index e7aca5ce5..2a116af05 100644
--- a/.github/workflows/nix_integration_tests.yaml
+++ b/.github/workflows/nix_integration_tests.yaml
@@ -30,13 +30,8 @@ jobs:
         USER: github_runner
     - name: Build
       run: nix develop .#test --command echo "Ok"
-    - name: Setup tmate session
-      uses: mxschmitt/action-tmate@v3
     - name: Python tests.
       run: |
-        sudo mkdir -p /run/opengl-driver/lib
-        sudo find /usr/lib . -name 'libcuda.so*' -exec ln -s {} /run/opengl-driver/lib/ \;
-        sudo find /usr/lib . -name 'libnvidia-ml.so*' -exec ln -s {} /run/opengl-driver/lib/ \;
         nix develop .#test --command pytest -svv integration-tests/ --release
       env:
         HF_TOKEN: ${{ secrets.HF_TOKEN }}
diff --git a/flake.nix b/flake.nix
index 674db0d82..45c28637d 100644
--- a/flake.nix
+++ b/flake.nix
@@ -95,6 +95,7 @@
                 cargo
                 rustfmt
                 clippy
+                cudaPackages.cudatoolkit
               ]
               ++ (with python3.pkgs; [
                 docker