fix the start gpu bug (deepjavalibrary#709)

Lokiiiiii · May 12, 2023 · a583888 · a583888
1 parent c8cbc0e
commit a583888
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/engines/python/src/main/java/ai/djl/python/engine/Connection.java b/engines/python/src/main/java/ai/djl/python/engine/Connection.java
@@ -189,7 +189,7 @@ private static String getVisibleDevices(int workerId, int tensorParallelDegree)
         // -> 0,2 and 3,7
         if (Utils.getenv("CUDA_VISIBLE_DEVICES") != null) {
             String[] devices = Utils.getenv("CUDA_VISIBLE_DEVICES").split(",");
-            sb.append(devices[tensorParallelDegree]);
+            sb.append(devices[workerId * tensorParallelDegree]);
             for (int i = 1; i < tensorParallelDegree; ++i) {
                 sb.append(',').append(devices[workerId * tensorParallelDegree + i]);
             }