diff --git a/config/ai.nix b/config/ai.nix
index ccb7bbd..00aa4a4 100644
--- a/config/ai.nix
+++ b/config/ai.nix
@@ -13,9 +13,35 @@
 
   users.users.mindforge = import ./user.nix;
 
+  # Setup drivers for NVIDIA GPU
+  services.xserver = {
+      enable = false;
+      videoDrivers = [ "nvidia" ];
+  };
+
+  hardware = {
+    nvidia = {
+      open = false;
+      modesetting.enable = true;
+      powerManagement.enable = false;
+      powerManagement.finegrained = false;
+      nvidiaSettings = true;
+    };
+
+    graphics = {
+      enable = true;
+      enable32Bit = true;
+    };
+  };
+
   services.ollama = {
     enable = true;
-    host = "0.0.0.0";
+    acceleration = "cuda";
+    package = pkgs.ollama-cuda.override {
+      cudaArches = [
+        "61"
+      ];
+    };
     openFirewall = true;
     environmentVariables = {
         OLLAMA_CONTEXT_LENGTH = "8192";
diff --git a/hardware/ai.nix b/hardware/ai.nix
index 75833fd..b176eb4 100644
--- a/hardware/ai.nix
+++ b/hardware/ai.nix
@@ -13,16 +13,16 @@
   boot.kernelModules = [ "kvm-intel" ];
   boot.extraModulePackages = [ ];
 
-  fileSystems."/" =
-    { device = "/dev/disk/by-uuid/b98eb7bb-f58c-4862-a234-1d72c9ff1187";
-      fsType = "ext4";
-    };
+  fileSystems."/" = {
+    device = "/dev/disk/by-partlabel/root";
+    fsType = "ext4";
+  };
 
-  fileSystems."/boot" =
-    { device = "/dev/disk/by-uuid/3D03-9579";
-      fsType = "vfat";
-      options = [ "fmask=0022" "dmask=0022" ];
-    };
+  fileSystems."/boot" = {
+    device = "/dev/disk/by-partlabel/EFI";
+    fsType = "vfat";
+    options = [ "fmask=0022" "dmask=0022" ];
+  };
 
   swapDevices = [ ];