diff --git a/test/compose/ai/docker-compose.yml b/test/compose/ai/docker-compose.yml
index e4162e1..7a6005c 100644
--- a/test/compose/ai/docker-compose.yml
+++ b/test/compose/ai/docker-compose.yml
@@ -3,6 +3,11 @@ volumes:
   open_webui: {}
   chromadb: {}
 
+networks:
+  host:
+    external: true
+    name: host
+
 services:
   open_webui:
     image: ghcr.io/open-webui/open-webui:ollama
@@ -27,15 +32,17 @@ services:
       - open_webui:/app/backend/data
     ports:
       - "8080:8080"
+    networks:
+      - host
     extra_hosts:
       - "host.docker.internal:host-gateway"
     deploy:
       resources:
         reservations:
           devices:
-            - driver: nvidia
-              count: 1
-              capabilities: [gpu]
+          - driver: cdi
+            device_ids: ['nvidia.com/gpu=all']
+            capabilities: [gpu]
 
   chromadb:
     hostname: chromadb
diff --git a/test/compose/nvidia-smi/docker-compose.yaml b/test/compose/nvidia-smi/docker-compose.yaml
index f61de85..f73a52a 100644
--- a/test/compose/nvidia-smi/docker-compose.yaml
+++ b/test/compose/nvidia-smi/docker-compose.yaml
@@ -1,11 +1,23 @@
 services:
-  test:
-    image: nvidia/cuda:12.3.1-base-ubuntu20.04
-    command: nvidia-smi
+  testgpupodman_count:
+    image: docker.io/ubuntu:latest
+    command: ["nvidia-smi"]
+    profiles: [gpu]
     deploy:
       resources:
         reservations:
           devices:
-            - driver: nvidia
-              count: 1
-              capabilities: [gpu]
\ No newline at end of file
+          - driver: nvidia
+            count: 1
+            capabilities: [gpu]
+  test:
+    image: docker.io/ubuntu:latest
+    command: ["nvidia-smi"]
+    deploy:
+      resources:
+        reservations:
+          devices:
+          - driver: cdi
+            count: 1
+            device_ids: ['nvidia.com/gpu=all']
+            capabilities: [gpu]
\ No newline at end of file