diff --git a/flake.nix b/flake.nix
index c510125..41cc300 100644
--- a/flake.nix
+++ b/flake.nix
@@ -67,5 +67,6 @@
           };
         };
       };
+      devenv = ./shell-modules/default.nix;
     };
 }
diff --git a/shell-modules/default.nix b/shell-modules/default.nix
new file mode 100644
index 0000000..91b74f7
--- /dev/null
+++ b/shell-modules/default.nix
@@ -0,0 +1,35 @@
+{ pkgs, ... }:
+
+pkgs.mkShell {
+  packages = with pkgs; [
+    (python312.withPackages (p: [
+      p.numpy
+      p.scikit-learn
+      p.scipy
+      p.pandas
+      p.matplotlib
+      p.torch-bin
+      # p.torchvision
+    ]))
+    libffi
+    openssl
+    stdenv.cc.cc
+    linuxPackages.nvidia_x11
+    binutils
+    cudatoolkit
+    libGLU
+    libGL
+  ];
+  LD_LIBRARY_PATH = pkgs.lib.makeLibraryPath (
+    with pkgs;
+    [
+      stdenv.cc.cc
+      linuxPackages.nvidia_x11
+      binutils
+      cudatoolkit
+      libGLU
+      libGL
+    ]
+  );
+  CUDA_PATH = pkgs.cudatoolkit;
+}