jaxlibWithCuda not using CUDA

Update:

Replacing the first line with { pkgs ? import <nixpkgs> { config.allowUnfree = true; config.cudaSupport = true } }: causes it to perform a lot of tests.

However, the same error persists, although now wandb must be removed since it fails some tests.