diff --git a/docker-compose-dev.yml b/docker-compose-dev.yml
index 1a38ed0..34b4d84 100644
--- a/docker-compose-dev.yml
+++ b/docker-compose-dev.yml
@@ -5,7 +5,6 @@ services:
       dockerfile: Dockerfile-dev
     image: kamco-cd-training-api:${IMAGE_TAG:-latest}
     container_name: kamco-cd-training-api
-    runtime: nvidia
     deploy:
       resources:
         reservations:
diff --git a/docker-compose-prod.yml b/docker-compose-prod.yml
index 4e796d8..a987358 100644
--- a/docker-compose-prod.yml
+++ b/docker-compose-prod.yml
@@ -5,6 +5,13 @@ services:
       dockerfile: Dockerfile
     image: kamco-train-api:${IMAGE_TAG:-latest}
     container_name: kamco-train-api
+    deploy:
+      resources:
+        reservations:
+          devices:
+            - driver: nvidia
+              count: all
+              capabilities: [gpu]
     expose:
       - "8080"
     environment: