With this shell.nix,

{ pkgs ? import <nixpkgs> {
  config.allowUnfree = true;
  config.cudaSupport = true;
} }:

pkgs.mkShell {
  packages = with pkgs; [
    (python3.withPackages (ps: [
      ps.torch
    ]))
  ];
}

Just a basic "is CUDA available" check fails.

$ nix-shell --run 'python'
Python 3.10.6 (main, Aug  1 2022, 20:38:21) [GCC 11.3.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import torch
>>> assert torch.cuda.is_available()
/nix/store/bf48f3zny7q08lg4hc4279fn3jw1lkpl-python3-3.10.6-env/lib/python3.10/site-packages/torch/cuda/__init__.py:83: UserWarning: CUDA initialization: CUDA unknown error - this may be due to an incorrectly set up environment, e.g. changing env variable CUDA_VISIBLE_DEVICES after program start. Setting the available devices to be zero. (Triggered internally at  /build/source/c10/cuda/CUDAFunctions.cpp:109.)
  return torch._C._cuda_getDeviceCount() > 0
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
AssertionError

(too late, but I'll still chime in with a comment on how I so far understand the landscape)

This is all that's needed from nixpkgs. The other requirements are imposed on the running system, and probably amount to having a /run/opengl-driver/lib/libcuda.so and some kernel module loaded. Both are deployed on NixOS when hardware.opengl.enable = true and the driver is nvidia

In reply to@ss:someonex.net

(too late, but I'll still chime in with a comment on how I so far understand the landscape)

This is all that's needed from nixpkgs. The other requirements are imposed on the running system, and probably amount to having a /run/opengl-driver/lib/libcuda.so and some kernel module loaded. Both are deployed on NixOS when hardware.opengl.enable = true and the driver is nvidia

FRidh: Hey, I just don't want to flood the nvidia-ml-py issue, so I post here: the thing is I don't know what you mean by "referring from the stub to the driver library"

I know that if you just run LD_PRELOAD=$(nix-build '<nixpkgs>' -A cudaPackages.cudatoolkit --no-out-link)/lib/stubs/libnvidia-ml.so nvidia-smi (or prepend the stub path in RUNPATH) without doing anything else, you're going to get a failure. If you meant something else (e.g. if you've heard of other ways to use the stubs than to bypass linker errors), then it might that I just lack knowledge