[inf2] Should not pass neuron device to PyTorch engine (deepjavalibra…

…ry#532)
Lokiiiiii · Mar 8, 2023 · 525dc58 · 525dc58
1 parent 8a90769
commit 525dc58
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/wlm/src/main/java/ai/djl/serving/wlm/ModelInfo.java b/wlm/src/main/java/ai/djl/serving/wlm/ModelInfo.java
@@ -208,7 +208,12 @@ public void load(Device device) throws ModelException, IOException {
  }
  }
  logger.info("Loading model {} on {}", id, device);
- builder.optDevice(device);
+ if ("nc".equals(device.getDeviceType()) && "PyTorch".equals(engineName)) {
+ // assume neuron only support PyTorch
+ logger.info("Bypass NC core allocation");
+ } else {
+ builder.optDevice(device);
+ }
 
  ZooModel<I, O> m = builder.build().loadModel();
  if (criteria != null) {