diff --git a/flake.nix b/flake.nix
index fabb8b0..249b75b 100644
--- a/flake.nix
+++ b/flake.nix
@@ -29,15 +29,17 @@
     packages."${system}" = {
       # just one main.cu file for now
       # compile it and place it to $out/bin/cuda-kernel
-      cuda-kernel = pkgs.stdenv.mkDerivation {
+      cuda-kernel = pkgs.cudaPackages.backendStdenv.mkDerivation {
         name = "cuda-kernel";
         src = ./.;
         buildInputs = with pkgs; [
           cudatoolkit
           cudaPackages.cuda_cudart
-          linuxPackages.nvidia_x11
+          linuxPackages.nvidia_x11_beta
         ];
-        LD_LIBRARY_PATH = "${pkgs.linuxPackages.nvidia_x11}/lib";
+        cudaSupport = true;
+        requiredSystemFeatures = [ "cuda" ];
+        LD_LIBRARY_PATH = "${pkgs.linuxPackages.nvidia_x11_beta}/lib";
         buildPhase = ''
           nvcc -o main main.cu
           mkdir -p $out/bin