diff --git a/modules/system/hardware/graphics/amdgpu.nix b/modules/system/hardware/graphics/amdgpu.nix
index 803a73c..c621d72 100644
--- a/modules/system/hardware/graphics/amdgpu.nix
+++ b/modules/system/hardware/graphics/amdgpu.nix
@@ -12,6 +12,10 @@ in
   };
   config = mkIf cfg.enable {
     boot.initrd.kernelModules = [ "amdgpu" ];
+
+    # Set acceleration to rocm
+    services.ollama.acceleration = "rocm";
+
     # Load amdgpu driver for Xorg and Wayland
     services.xserver = {
       enable = true;
diff --git a/modules/system/hardware/graphics/nvidia.nix b/modules/system/hardware/graphics/nvidia.nix
index 0641a01..2f229b4 100644
--- a/modules/system/hardware/graphics/nvidia.nix
+++ b/modules/system/hardware/graphics/nvidia.nix
@@ -14,6 +14,9 @@ in
     # Load nvidia driver for Xorg and Wayland
     services.xserver.videoDrivers = ["nvidia"];
 
+    # Set acceleration to cuda
+    services.ollama.acceleration = "cuda";
+
     hardware = {
       graphics = {
         enable = true;
diff --git a/modules/system/services/ollama.nix b/modules/system/services/ollama.nix
index fc5fc65..1dd8e81 100644
--- a/modules/system/services/ollama.nix
+++ b/modules/system/services/ollama.nix
@@ -3,7 +3,6 @@
 with lib;
 let
   cfg = config.modules.system.services.ollama;
-  nvidiaEnabled = config.modules.system.hardware.nvidia.enable;
 in
 {
   options.modules.system.services.ollama = {
@@ -14,8 +13,6 @@ in
   config = mkIf cfg.enable {
     services.ollama = {
       enable = true;
-
-      acceleration = if nvidiaEnabled then "cuda" else null;
     };
   };
 }