Set specific gpu 0 as well

triton-inference-server · Jul 24, 2024 · 64d5b90 · 64d5b90
1 parent ea043eb
commit 64d5b90
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/python/test/test_api.py b/python/test/test_api.py
@@ -357,8 +357,11 @@ def test_stop(self):
                     {
                         "backend": "python",
                         "parameters": {"decoupled": {"string_value": "False"}},
-                        # Keep instance count low for fast startup/cleanup
-                        "instance_group": [{"kind": "KIND_GPU", "count": 1}],
+                        # Keep instance count low for fast startup/cleanup.
+                        # Alternatively can use KIND_CPU here, but keeping gpus/count explicit.
+                        "instance_group": [
+                            {"kind": "KIND_GPU", "gpus": [0], "count": 1}
+                        ],
                     }
                 )
             },