[canary] Fixes canary test (#505)

deepjavalibrary · Sep 4, 2024 · a4fdcab · a4fdcab
1 parent d084f7e
commit a4fdcab
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/canary/src/main/java/ai/djl/canary/CanaryTest.java b/canary/src/main/java/ai/djl/canary/CanaryTest.java
@@ -272,7 +272,7 @@ private static void testXgboost() throws ModelException, IOException, TranslateE
         try (Model model = Model.newInstance("canary", "XGBoost")) {
             model.load(modelDir, "regression");
             Predictor<NDList, NDList> predictor = model.newPredictor(new NoopTranslator());
-            try (NDManager manager = NDManager.newBaseManager()) {
+            try (NDManager manager = NDManager.newBaseManager("XGBoost")) {
                 NDArray array = manager.ones(new Shape(10, 13));
                 NDList output = predictor.predict(new NDList(array));
                 float[] result = output.singletonOrThrow().toFloatArray();
@@ -307,7 +307,7 @@ private static void testLightgbm() throws ModelException, IOException, Translate
 
         try (ZooModel<NDList, NDList> model = criteria.loadModel();
                 Predictor<NDList, NDList> predictor = model.newPredictor()) {
-            try (NDManager manager = NDManager.newBaseManager()) {
+            try (NDManager manager = NDManager.newBaseManager("LightGBM")) {
                 NDArray array = manager.ones(new Shape(10, 4));
                 NDList output = predictor.predict(new NDList(array));
                 float[] result = output.singletonOrThrow().toFloatArray();