diff --git a/flake.nix b/flake.nix
index 41cc300..bd2278c 100644
--- a/flake.nix
+++ b/flake.nix
@@ -17,10 +17,12 @@
       url = "gitlab:doronbehar/nix-matlab";
       inputs.nixpkgs.follows = "nixpkgs";
     };
+    flake-utils.url = "github:numtide/flake-utils";
   };
 
   outputs =
     {
+      self,
       nixpkgs,
       home-manager,
       stylix,
@@ -67,6 +69,6 @@
           };
         };
       };
-      devenv = ./shell-modules/default.nix;
+      lib = import ./shell-modules/default.nix self.inputs;
     };
 }
diff --git a/shell-modules/default.nix b/shell-modules/default.nix
index 91b74f7..d08cac4 100644
--- a/shell-modules/default.nix
+++ b/shell-modules/default.nix
@@ -1,35 +1,30 @@
-{ pkgs, ... }:
+{ nixpkgs, flake-utils, ... }:
 
-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
+let
+  imports = [
+    ./languags/python.nix
   ];
-  LD_LIBRARY_PATH = pkgs.lib.makeLibraryPath (
-    with pkgs;
-    [
-      stdenv.cc.cc
-      linuxPackages.nvidia_x11
-      binutils
-      cudatoolkit
-      libGLU
-      libGL
-    ]
-  );
-  CUDA_PATH = pkgs.cudatoolkit;
+in
+{
+
+  mkShell =
+    attrs:
+    (flake-utils.lib.eachDefaultSystem (
+      system:
+      let
+        pkgs = import nixpkgs {
+          inherit system;
+          config.allowUnfree = true;
+        };
+        modules = [
+          attrs
+        ] ++ imports;
+        evaluated = nixpkgs.lib.evalModules { inherit modules; };
+      in
+      {
+        devShells.default = pkgs.mkShell {
+          TEST_ENV = builtins.trace evaluated.config "HELLO";
+        };
+      }
+    ));
 }
diff --git a/shell-modules/languages/python.nix b/shell-modules/languages/python.nix
new file mode 100644
index 0000000..e69de29