diff --git a/hosts/default.nix b/hosts/default.nix
index 14d611a..99c0c85 100644
--- a/hosts/default.nix
+++ b/hosts/default.nix
@@ -20,13 +20,6 @@
 
     perClass = class: {
       modules = ["${self}/modules/${class}/default.nix"];
-
-      # modules = with inputs;
-      #   (nixpkgs.lib.optionals (class == "nixos") [
-      #     niri-flake.nixosModules.niri
-      #   ])
-      #   ++ (nixpkgs.lib.optionals (class == "darwin") [
-      #   ]);
     };
 
     hosts = {
diff --git a/modules/nixos/shell/niri.nix b/modules/nixos/shell/niri.nix
index 761c597..e4d8674 100644
--- a/modules/nixos/shell/niri.nix
+++ b/modules/nixos/shell/niri.nix
@@ -1,12 +1,15 @@
 {
-  lib,
   config,
+  lib,
+  inputs,
   ...
 }: let
   cfg = config.marleyos.shell.niri;
 in {
   options.marleyos.shell.niri.enable = lib.mkEnableOption "niri";
 
+  imports = [inputs.niri-flake.nixosModules.niri];
+
   config = lib.mkIf cfg.enable {
     programs.niri = {
       enable = true;