[inf2] Fixes NeuronUtils.getNeuronCores() bug (deepjavalibrary#531)

Lokiiiiii · Mar 8, 2023 · 8a90769 · 8a90769
1 parent c6146eb
commit 8a90769
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/serving/src/main/java/ai/djl/serving/util/NeuronUtils.java b/serving/src/main/java/ai/djl/serving/util/NeuronUtils.java
@@ -67,14 +67,14 @@ public static int getNeuronCores() {
  if (!hasNeuron()) {
  return 0;
  }
- try (Stream<Path> paths = Files.walk(Paths.get("/dev"))) {
- int nd = (int) paths.filter(ele -> ele.startsWith("neuron")).count();
+ try (Stream<Path> paths = Files.list(Paths.get("/dev"))) {
+ long nd = paths.filter(p -> p.getFileName().toString().startsWith("neuron")).count();
  if (isInf1()) {
  // inf1 has 4 cores on each device
- return nd * 4;
+ return (int) nd * 4;
  }
  // inf2 has 2 cores on each device
- return nd * 2;
+ return (int) nd * 2;
  } catch (IOException e) {
  throw new AssertionError("Failed to list neuron cores", e);
  }