[api] Tries to use the same device for alternative NDManager (#3146)

deepjavalibrary · May 1, 2024 · 00c78f6 · 00c78f6
1 parent 1d3613a
commit 00c78f6
Showing 1 changed file with 8 additions and 2 deletions.
diff --git a/api/src/main/java/ai/djl/ndarray/BaseNDManager.java b/api/src/main/java/ai/djl/ndarray/BaseNDManager.java
@@ -14,6 +14,7 @@
 
 import ai.djl.Device;
 import ai.djl.engine.Engine;
+import ai.djl.engine.StandardCapabilities;
 import ai.djl.ndarray.types.DataType;
 import ai.djl.ndarray.types.Shape;
 import ai.djl.util.PairList;
@@ -62,8 +63,13 @@ protected BaseNDManager(NDManager parent, Device device) {
         uid = UUID.randomUUID().toString();
         Engine engine = getEngine().getAlternativeEngine();
         if (engine != null) {
-            // Use the default device
-            alternativeManager = engine.newBaseManager();
+            // Use the same device if possible for efficiency
+            if (this.device.isGpu() && engine.hasCapability(StandardCapabilities.CUDA)) {
+                alternativeManager = engine.newBaseManager(this.device);
+            } else {
+                // Use the default device
+                alternativeManager = engine.newBaseManager();
+            }
         }
     }