diff --git a/api/src/main/java/ai/djl/util/NeuronUtils.java b/api/src/main/java/ai/djl/util/NeuronUtils.java index 685dbd77251..762b4e1469e 100644 --- a/api/src/main/java/ai/djl/util/NeuronUtils.java +++ b/api/src/main/java/ai/djl/util/NeuronUtils.java @@ -16,6 +16,7 @@ import org.slf4j.LoggerFactory; import java.io.IOException; +import java.io.InputStream; import java.nio.file.Files; import java.nio.file.Path; import java.nio.file.Paths; @@ -86,6 +87,14 @@ public static int getNeuronCoresForDevice(String location) { if (!Files.exists(path)) { return 0; } + Path file = path.resolve("core_count"); + if (Files.exists(file)) { + try (InputStream is = Files.newInputStream(file)) { + return Integer.parseInt(Utils.toString(is)); + } catch (IOException e) { + throw new AssertionError("Failed to read core_count file", e); + } + } int count = 0; try (Stream dev = Files.list(path)) { return Math.toIntExact(dev.filter(p -> matches(p, "neuron_core")).count());