Add CUDA libs and deps

modules/overlays.nix

@@ -6,7 +6,18 @@
 {
   nixpkgs.config.allowUnfreePredicate =
     pkg:
-    builtins.elem (nixpkgs.lib.getName pkg) [
+    nixpkgs.lib.all (
+      license:
+      license.free
+      || nixpkgs.lib.elem license.shortName [
+        "CUDA EULA"
+        "cuDNN EULA"
+        "cuSPARSELt EULA"
+        "cuTENSOR EULA"
+        "NVidia OptiX EULA"
+      ]
+    ) (nixpkgs.lib.toList pkg.meta.license)
+    || builtins.elem (nixpkgs.lib.getName pkg) [
       "nvidia-x11"
     ];
 

modules/system/options/nixos/graphics/cuda.nix

@@ -1,4 +1,9 @@
-{ lib, config, ... }:
+{
+  pkgs,
+  lib,
+  config,
+  ...
+}:
 with lib;
 let
   cfg = config.os.virtualisation.cuda;
@@ -10,6 +15,22 @@ in
   config = mkIf cfg.enable {
 
     hardware.nvidia-container-toolkit.enable = true;
+    environment.systemPackages = with pkgs.cudaPackages; [
+      cudatoolkit
+      cudnn
+      libcublas
+      libcufile
+      libcusparse
+      nccl
+    ];
+
+    environment.variables = {
+      CUDA_PATH = "${pkgs.cudatoolkit}";
+      LD_LIBRARY_PATH = "${pkgs.linuxPackages.nvidia_x11}/lib:$LD_LIBRARY_PATH";
+      EXTRA_LDFLAGS = "-L/lib -L${pkgs.linuxPackages.nvidia_x11}/lib";
+      EXTRA_CCFLAGS = "-I/usr/include";
+    };
 
   };
+
 }