diff --git a/flake.nix b/flake.nix
index 00e178a..e032fcf 100644
--- a/flake.nix
+++ b/flake.nix
@@ -5,6 +5,13 @@
     flake-utils.url = "github:numtide/flake-utils";
   };
 
+  nixConfig = {
+    extra-substituters = [ "https://cuda-maintainers.cachix.org" ];
+    extra-trusted-public-keys = [
+      "cuda-maintainers.cachix.org-1:0dq3bujKpuEPMCX6U4WylrUDZ9JyUG0VpVZa7CNfq5E="
+    ];
+  };
+
   outputs = { self, nixpkgs, flake-utils }:
     flake-utils.lib.eachDefaultSystem (system:
       let
@@ -64,7 +71,7 @@
             fetchSubmodules = true;
           };
 
-          propagatedBuildInputs = with pythonPackages; [ numpy torch ];
+          propagatedBuildInputs = with pythonPackages; [ numpy torchWithCuda ];
 
           nativeBuildInputs = with pythonPackages; [
             setuptools
@@ -100,7 +107,7 @@
           propagatedBuildInputs = with pythonPackages; [
             requests
             pyyaml
-            torch
+            torchWithCuda
             tqdm
             psutil
             torchvision
@@ -154,7 +161,7 @@
             numpy
             pandas
             pyyaml
-            torch
+            torchWithCuda
             yacs
             torchvision
             ftfy
@@ -194,7 +201,7 @@
           python312Packages.opencv4 = custom-opencv4;
         };
         devShells.default = pkgs.mkShell {
-          buildInputs = [ (pkgs.python312.withPackages (ps: [ ])) ];
+          buildInputs = [ (pkgs.python312.withPackages (ps: [ ultralytics ])) ];
         };
       });
 }