diff --git a/docker-compose.yml b/docker-compose.yml
index f93ff34..dfd6bc6 100644
--- a/docker-compose.yml
+++ b/docker-compose.yml
@@ -11,9 +11,9 @@ x-base_service: &base_service
       resources:
         reservations:
           devices:
-              - driver: nvidia
-                device_ids: ['0']
-                capabilities: [gpu]
+            - driver: nvidia
+              count: all
+              capabilities: [gpu]
 volumes:
   data:
   output: