diff --git a/README.md b/README.md
index b363e01..23d3349 100644
--- a/README.md
+++ b/README.md
@@ -31,7 +31,7 @@ Instead of `vm.nix`, `nixos-shell` also accepts other modules on the command lin
 $ nixos-shell some-nix-module.nix
 ```
 
-You can also start a vm from a flake's `nixosConfigurations` using the `--flake` flag.
+You can also start a vm from a flake's `nixosConfigurations` or `nixosModules` output using the `--flake` flag.
 
 ```console
 $ nixos-shell --flake github:Mic92/nixos-shell#vm-forward
@@ -52,6 +52,16 @@ This will run the `vm-forward` example.
 >}
 >```
 
+When using the `--flake` flag, if no attribute is given, `nixos-shell` tries the following flake output attributes:
+- `packages.<system>.nixosConfigurations.<vm>`
+- `nixosConfigurations.<vm>`
+- `nixosModules.<vm>`
+
+If an attribute _name_ is given, `nixos-shell` tries the following flake output attributes:
+- `packages.<system>.nixosConfigurations.<name>`
+- `nixosConfigurations.<name>`
+- `nixosModules.<name>`
+
 ## Terminating the virtual machine
 
 Type `Ctrl-a x` to exit the virtual machine.
diff --git a/share/nixos-shell.nix b/share/nixos-shell.nix
index a498ff3..65380e0 100644
--- a/share/nixos-shell.nix
+++ b/share/nixos-shell.nix
@@ -6,31 +6,49 @@
 , flakeAttr ? null
 }:
 let
+  lib = (import nixpkgs { }).lib;
+
   nixos-shell = import ./modules/nixos-shell.nix;
   nixos-shell-config = import ./modules/nixos-shell-config.nix;
 
-  flake = builtins.getFlake flakeUri;
-  flakeSystem = flake.outputs.packages."${system}".nixosConfigurations."${flakeAttr}" or flake.outputs.nixosConfigurations."${flakeAttr}";
-in
-if flakeUri != null then
-  flakeSystem.override
-    (attrs: {
-      modules =
-        let
-          nixosShellModules =
-            if flakeSystem ? options.nixos-shell then
-              [ nixos-shell-config ]
-            else
-              [ nixos-shell nixos-shell-config ];
-        in
-        attrs.modules ++ nixosShellModules;
-    })
-else
-  import "${toString nixpkgs}/nixos/lib/eval-config.nix" {
+  defaultTo = default: e: if e == null then default else e;
+
+  getFlakeOutput = path: lib.attrByPath path null flake.outputs;
+
+  mkShellSystem = config: import "${toString nixpkgs}/nixos/lib/eval-config.nix" {
     inherit system;
     modules = [
-      configuration
+      config
       nixos-shell
       nixos-shell-config
     ];
-  }
+  };
+
+  flake = builtins.getFlake flakeUri;
+
+  flakeSystem = defaultTo
+    (getFlakeOutput [ "nixosConfigurations" "${flakeAttr}" ])
+    (getFlakeOutput [ "packages" "${system}" "nixosConfigurations" "${flakeAttr}" ]);
+
+  flakeModule = getFlakeOutput [ "nixosModules" "${flakeAttr}" ];
+in
+if flakeUri != null then
+  if flakeSystem != null then
+    flakeSystem.override
+      (attrs: {
+        modules =
+          let
+            nixosShellModules =
+              if flakeSystem ? options.nixos-shell then
+                [ nixos-shell-config ]
+              else
+                [ nixos-shell nixos-shell-config ];
+          in
+          attrs.modules ++ nixosShellModules;
+      })
+  else if flakeModule != null then
+    mkShellSystem flakeModule
+  else
+    throw "cannot find flake attribute '${flakeUri}#${flakeAttr}'"
+else
+  mkShellSystem configuration