Repository: deepjavalibrary/djl
Pull request #3297 (2024-07-03 15:27:33)

[api] Refactor nms for yolo translator
## Description ##

Brief description of what this PR is about

- If this change is a backward incompatible change, why must this change be made?
- Interesting edge cases to note here



Patch:
diff --git a/api/src/main/java/ai/djl/modality/cv/BufferedImageFactory.java b/api/src/main/java/ai/djl/modality/cv/BufferedImageFactory.java
index 40735ddec..d18b6922b 100644
--- a/api/src/main/java/ai/djl/modality/cv/BufferedImageFactory.java
+++ b/api/src/main/java/ai/djl/modality/cv/BufferedImageFactory.java
@@ -403,28 +403,38 @@ public class BufferedImageFactory extends ImageFactory {
             float b = RandomUtils.nextFloat();
             int imageWidth = image.getWidth();
             int imageHeight = image.getHeight();
-            int x = (int) (mask.getX() * imageWidth);
-            int y = (int) (mask.getY() * imageHeight);
-            float[][] probDist = mask.getProbDist();
-            // Correct some coordinates of box when going out of image
-            if (x < 0) {
-                x = 0;
-            }
-            if (y < 0) {
-                y = 0;
+            int x = 0;
+            int y = 0;
+            int w = imageWidth;
+            int h = imageHeight;
+            if (!mask.isFullImageMask()) {
+                x = (int) (mask.getX() * imageWidth);
+                y = (int) (mask.getY() * imageHeight);
+                w = (int) (mask.getWidth() * imageWidth);
+                h = (int) (mask.getHeight() * imageHeight);
+                // Correct some coordinates of box when going out of image
+                if (x < 0) {
+                    x = 0;
+                }
+                if (y < 0) {
+                    y = 0;
+                }
             }
+            float[][] probDist = mask.getProbDist();
 
             BufferedImage maskImage =
                     new BufferedImage(
-                            probDist.length, probDist[0].length, BufferedImage.TYPE_INT_ARGB);
-            for (int xCor = 0; xCor < probDist.length; xCor++) {
-                for (int yCor = 0; yCor < probDist[xCor].length; yCor++) {
-                    float opacity = probDist[xCor][yCor] * 0.8f;
+                            probDist[0].length, probDist.length, BufferedImage.TYPE_INT_ARGB);
+            for (int yCor = 0; yCor < probDist.length; yCor++) {
+                for (int xCor = 0; xCor < probDist[0].length; xCor++) {
+                    float opacity = probDist[yCor][xCor] * 0.8f;
                     maskImage.setRGB(xCor, yCor, new Color(r, g, b, opacity).darker().getRGB());
                 }
             }
+            java.awt.Image scaled = maskImage.getScaledInstance(w, h, java.awt.Image.SCALE_SMOOTH);
+
             Graphics2D gR = (Graphics2D) image.getGraphics();
-            gR.drawImage(maskImage, x, y, null);
+            gR.drawImage(scaled, x, y, null);
             gR.dispose();
         }
 
diff --git a/api/src/main/java/ai/djl/modality/cv/output/Mask.java b/api/src/main/java/ai/djl/modality/cv/output/Mask.java
index 0c4a9fd66..e9ada0710 100644
--- a/api/src/main/java/ai/djl/modality/cv/output/Mask.java
+++ b/api/src/main/java/ai/djl/modality/cv/output/Mask.java
@@ -21,6 +21,7 @@ public class Mask extends Rectangle {
 
     private static final long serialVersionUID = 1L;
     private float[][] probDist;
+    private boolean fullImageMask;
 
     /**
      * Constructs a Mask with the given data.
@@ -32,8 +33,29 @@ public class Mask extends Rectangle {
      * @param dist the probability distribution for each pixel in the rectangle
      */
     public Mask(double x, double y, double width, double height, float[][] dist) {
+        this(x, y, width, height, dist, false);
+    }
+
+    /**
+     * Constructs a Mask with the given data.
+     *
+     * @param x the left coordinate of the bounding rectangle
+     * @param y the top coordinate of the bounding rectangle
+     * @param width the width of the bounding rectangle
+     * @param height the height of the bounding rectangle
+     * @param dist the probability distribution for each pixel in the rectangle
+     * @param fullImageMask if the mask if for full image
+     */
+    public Mask(
+            double x,
+            double y,
+            double width,
+            double height,
+            float[][] dist,
+            boolean fullImageMask) {
         super(x, y, width, height);
         this.probDist = dist;
+        this.fullImageMask = fullImageMask;
     }
 
     /**
@@ -44,4 +66,13 @@ public class Mask extends Rectangle {
     public float[][] getProbDist() {
         return probDist;
     }
+
+    /**
+     * Returns if the mask is for full image.
+     *
+     * @return if the mask is for full image
+     */
+    public boolean isFullImageMask() {
+        return fullImageMask;
+    }
 }
diff --git a/api/src/main/java/ai/djl/modality/cv/output/Rectangle.java b/api/src/main/java/ai/djl/modality/cv/output/Rectangle.java
index 92afc6032..b7b16127a 100644
--- a/api/src/main/java/ai/djl/modality/cv/output/Rectangle.java
+++ b/api/src/main/java/ai/djl/modality/cv/output/Rectangle.java
@@ -14,6 +14,7 @@ package ai.djl.modality.cv.output;
 
 import java.util.ArrayList;
 import java.util.List;
+import java.util.PriorityQueue;
 
 /**
  * A {@code Rectangle} specifies an area in a coordinate space that is enclosed by the {@code
@@ -152,4 +153,81 @@ public class Rectangle implements BoundingBox {
         return String.format(
                 "{\"x\"=%.3f, \"y\"=%.3f, \"width\"=%.3f, \"height\"=%.3f}", x, y, width, height);
     }
+
+    /**
+     * Applies nms (non-maximum suppression) to the list of rectangles.
+     *
+     * @param boxes an list of {@code Rectangle}
+     * @param scores a list of scores
+     * @param nmsThreshold the nms threshold
+     * @return the filtered list with the index of the original list
+     */
+    public static List<Integer> nms(
+            List<Rectangle> boxes, List<Double> scores, float nmsThreshold) {
+        List<Integer> ret = new ArrayList<>();
+        PriorityQueue<Integer> pq =
+                new PriorityQueue<>(
+                        50,
+                        (lhs, rhs) -> {
+                            // Intentionally reversed to put high confidence at the head of the
+                            // queue.
+                            return Double.compare(scores.get(rhs), scores.get(lhs));
+                        });
+        for (int i = 0; i < boxes.size(); ++i) {
+            pq.add(i);
+        }
+
+        // do non maximum suppression
+        while (!pq.isEmpty()) {
+            // insert detection with max confidence
+            int[] detections = pq.stream().mapToInt(Integer::intValue).toArray();
+            ret.add(detections[0]);
+            Rectangle box = boxes.get(detections[0]);
+            pq.clear();
+            for (int i = 1; i < detections.length; i++) {
+                int detection = detections[i];
+                Rectangle location = boxes.get(detection);
+                if (box.boxIou(location) < nmsThreshold) {
+                    pq.add(detection);
+                }
+            }
+        }
+        return ret;
+    }
+
+    private double boxIou(Rectangle other) {
+        double intersection = intersection(other);
+        double union =
+                getWidth() * getHeight() + other.getWidth() * other.getHeight() - intersection;
+        return intersection / union;
+    }
+
+    private double intersection(Rectangle b) {
+        double w =
+                overlap(
+                        (getX() * 2 + getWidth()) / 2,
+                        getWidth(),
+                        (b.getX() * 2 + b.getWidth()) / 2,
+                        b.getWidth());
+        double h =
+                overlap(
+                        (getY() * 2 + getHeight()) / 2,
+                        getHeight(),
+                        (b.getY() * 2 + b.getHeight()) / 2,
+                        b.getHeight());
+        if (w < 0 || h < 0) {
+            return 0;
+        }
+        return w * h;
+    }
+
+    private double overlap(double x1, double w1, double x2, double w2) {
+        double l1 = x1 - w1 / 2;
+        double l2 = x2 - w2 / 2;
+        double left = Math.max(l1, l2);
+        double r1 = x1 + w1 / 2;
+        double r2 = x2 + w2 / 2;
+        double right = Math.min(r1, r2);
+        return right - left;
+    }
 }
diff --git a/api/src/main/java/ai/djl/modality/cv/translator/InstanceSegmentationTranslator.java b/api/src/main/java/ai/djl/modality/cv/translator/InstanceSegmentationTranslator.java
index efd05d7be..458122bbc 100644
--- a/api/src/main/java/ai/djl/modality/cv/translator/InstanceSegmentationTranslator.java
+++ b/api/src/main/java/ai/djl/modality/cv/translator/InstanceSegmentationTranslator.java
@@ -12,7 +12,6 @@
  */
 package ai.djl.modality.cv.translator;
 
-import ai.djl.modality.cv.Image;
 import ai.djl.modality.cv.output.BoundingBox;
 import ai.djl.modality.cv.output.DetectedObjects;
 import ai.djl.modality.cv.output.Mask;
@@ -68,14 +67,6 @@ public class InstanceSegmentationTranslator extends BaseImageTranslator<Detected
         }
     }
 
-    /** {@inheritDoc} */
-    @Override
-    public NDList processInput(TranslatorContext ctx, Image image) {
-        ctx.setAttachment("originalHeight", image.getHeight());
-        ctx.setAttachment("originalWidth", image.getWidth());
-        return super.processInput(ctx, image);
-    }
-
     /** {@inheritDoc} */
     @Override
     public DetectedObjects processOutput(TranslatorContext ctx, NDList list) {
@@ -102,18 +93,15 @@ public class InstanceSegmentationTranslator extends BaseImageTranslator<Detected
                 double w = box[2] / rescaledWidth - x;
                 double h = box[3] / rescaledHeight - y;
 
-                int maskW = (int) (w * (int) ctx.getAttachment("originalWidth"));
-                int maskH = (int) (h * (int) ctx.getAttachment("originalHeight"));
-
                 // Reshape mask to actual image bounding box shape.
                 NDArray array = masks.get(i);
                 Shape maskShape = array.getShape();
-                array = array.reshape(maskShape.addAll(new Shape(1)));
-                NDArray maskArray = NDImageUtils.resize(array, maskW, maskH).transpose();
-                float[] flattened = maskArray.toFloatArray();
-                float[][] maskFloat = new float[maskW][maskH];
-                for (int j = 0; j < maskW; j++) {
-                    System.arraycopy(flattened, j * maskH, maskFloat[j], 0, maskH);
+                int maskH = (int) maskShape.get(0);
+                int maskW = (int) maskShape.get(1);
+                float[] flattened = array.toFloatArray();
+                float[][] maskFloat = new float[maskH][maskW];
+                for (int j = 0; j < maskH; j++) {
+                    System.arraycopy(flattened, j * maskW, maskFloat[j], 0, maskW);
                 }
                 Mask mask = new Mask(x, y, w, h, maskFloat);
 
diff --git a/api/src/main/java/ai/djl/modality/cv/translator/YoloV5Translator.java b/api/src/main/java/ai/djl/modality/cv/translator/YoloV5Translator.java
index c31353766..67682d7c8 100644
--- a/api/src/main/java/ai/djl/modality/cv/translator/YoloV5Translator.java
+++ b/api/src/main/java/ai/djl/modality/cv/translator/YoloV5Translator.java
@@ -24,7 +24,6 @@ import java.util.ArrayList;
 import java.util.List;
 import java.util.Locale;
 import java.util.Map;
-import java.util.PriorityQueue;
 
 /**
  * A translator for YoloV5 models. This was tested with ONNX exported Yolo models. For details check
@@ -68,104 +67,58 @@ public class YoloV5Translator extends ObjectDetectionTranslator {
         return builder;
     }
 
-    protected double boxIntersection(Rectangle a, Rectangle b) {
-        double w =
-                overlap(
-                        (a.getX() * 2 + a.getWidth()) / 2,
-                        a.getWidth(),
-                        (b.getX() * 2 + b.getWidth()) / 2,
-                        b.getWidth());
-        double h =
-                overlap(
-                        (a.getY() * 2 + a.getHeight()) / 2,
-                        a.getHeight(),
-                        (b.getY() * 2 + b.getHeight()) / 2,
-                        b.getHeight());
-        if (w < 0 || h < 0) {
-            return 0;
-        }
-        return w * h;
-    }
-
-    protected double boxIou(Rectangle a, Rectangle b) {
-        return boxIntersection(a, b) / boxUnion(a, b);
-    }
-
-    protected double boxUnion(Rectangle a, Rectangle b) {
-        double i = boxIntersection(a, b);
-        return (a.getWidth()) * (a.getHeight()) + (b.getWidth()) * (b.getHeight()) - i;
-    }
-
-    protected DetectedObjects nms(List<IntermediateResult> list) {
+    protected DetectedObjects nms(
+            List<Rectangle> boxes, List<Integer> classIds, List<Float> scores) {
         List<String> retClasses = new ArrayList<>();
         List<Double> retProbs = new ArrayList<>();
         List<BoundingBox> retBB = new ArrayList<>();
 
-        for (int k = 0; k < classes.size(); k++) {
-            // 1.find max confidence per class
-            PriorityQueue<IntermediateResult> pq =
-                    new PriorityQueue<>(
-                            50,
-                            (lhs, rhs) -> {
-                                // Intentionally reversed to put high confidence at the head of the
-                                // queue.
-                                return Double.compare(rhs.getConfidence(), lhs.getConfidence());
-                            });
-
-            for (IntermediateResult intermediateResult : list) {
-                if (intermediateResult.getDetectedClass() == k) {
-                    pq.add(intermediateResult);
+        for (int classId = 0; classId < classes.size(); classId++) {
+            List<Rectangle> r = new ArrayList<>();
+            List<Double> s = new ArrayList<>();
+            List<Integer> map = new ArrayList<>();
+            for (int j = 0; j < classIds.size(); ++j) {
+                if (classIds.get(j) == classId) {
+                    r.add(boxes.get(j));
+                    s.add(scores.get(j).doubleValue());
+                    map.add(j);
                 }
             }
-
-            // 2.do non maximum suppression
-            while (pq.size() > 0) {
-                // insert detection with max confidence
-                IntermediateResult[] a = new IntermediateResult[pq.size()];
-                IntermediateResult[] detections = pq.toArray(a);
-                Rectangle rec = detections[0].getLocation();
-                retClasses.add(detections[0].id);
-                retProbs.add(detections[0].confidence);
+            if (r.isEmpty()) {
+                continue;
+            }
+            List<Integer> nms = Rectangle.nms(r, s, nmsThreshold);
+            for (int index : nms) {
+                int pos = map.get(index);
+                int id = classIds.get(pos);
+                retClasses.add(classes.get(id));
+                retProbs.add(scores.get(pos).doubleValue());
+                Rectangle rect = boxes.get(pos);
                 if (applyRatio) {
                     retBB.add(
                             new Rectangle(
-                                    rec.getX() / imageWidth,
-                                    rec.getY() / imageHeight,
-                                    rec.getWidth() / imageWidth,
-                                    rec.getHeight() / imageHeight));
+                                    rect.getX() / imageWidth,
+                                    rect.getY() / imageHeight,
+                                    rect.getWidth() / imageWidth,
+                                    rect.getHeight() / imageHeight));
                 } else {
-                    retBB.add(
-                            new Rectangle(rec.getX(), rec.getY(), rec.getWidth(), rec.getHeight()));
-                }
-                pq.clear();
-                for (int j = 1; j < detections.length; j++) {
-                    IntermediateResult detection = detections[j];
-                    Rectangle location = detection.getLocation();
-                    if (boxIou(rec, location) < nmsThreshold) {
-                        pq.add(detection);
-                    }
+                    retBB.add(rect);
                 }
             }
         }
         return new DetectedObjects(retClasses, retProbs, retBB);
     }
 
-    protected double overlap(double x1, double w1, double x2, double w2) {
-        double l1 = x1 - w1 / 2;
-        double l2 = x2 - w2 / 2;
-        double left = Math.max(l1, l2);
-        double r1 = x1 + w1 / 2;
-        double r2 = x2 + w2 / 2;
-        double right = Math.min(r1, r2);
-        return right - left;
-    }
-
     protected DetectedObjects processFromBoxOutput(NDList list) {
         float[] flattened = list.get(0).toFloatArray();
-        ArrayList<IntermediateResult> intermediateResults = new ArrayList<>();
         int sizeClasses = classes.size();
         int stride = 5 + sizeClasses;
         int size = flattened.length / stride;
+
+        ArrayList<Rectangle> boxes = new ArrayList<>();
+        ArrayList<Float> scores = new ArrayList<>();
+        ArrayList<Integer> classIds = new ArrayList<>();
+
         for (int i = 0; i < size; i++) {
             int indexBase = i * stride;
             float maxClass = 0;
@@ -184,11 +137,12 @@ public class YoloV5Translator extends ObjectDetectionTranslator {
                 float h = flattened[indexBase + 3];
                 Rectangle rect =
                         new Rectangle(Math.max(0, xPos - w / 2), Math.max(0, yPos - h / 2), w, h);
-                intermediateResults.add(
-                        new IntermediateResult(classes.get(maxIndex), score, maxIndex, rect));
+                boxes.add(rect);
+                scores.add(score);
+                classIds.add(maxIndex);
             }
         }
-        return nms(intermediateResults);
+        return nms(boxes, classIds, scores);
     }
 
     private DetectedObjects processFromDetectOutput() {
@@ -279,49 +233,4 @@ public class YoloV5Translator extends ObjectDetectionTranslator {
             return new YoloV5Translator(this);
         }
     }
-
-    protected static final class IntermediateResult {
-
-        /**
-         * A sortable score for how good the recognition is relative to others. Higher should be
-         * better.
-         */
-        private double confidence;
-
-        /** Display name for the recognition. */
-        private int detectedClass;
-
-        /**
-         * A unique identifier for what has been recognized. Specific to the class, not the instance
-         * of the object.
-         */
-        private String id;
-
-        /** Optional location within the source image for the location of the recognized object. */
-        private Rectangle location;
-
-        IntermediateResult(String id, double confidence, int detectedClass, Rectangle location) {
-            this.confidence = confidence;
-            this.id = id;
-            this.detectedClass = detectedClass;
-            this.location = location;
-        }
-
-        public double getConfidence() {
-            return confidence;
-        }
-
-        public int getDetectedClass() {
-            return detectedClass;
-        }
-
-        public String getId() {
-            return id;
-        }
-
-        public Rectangle getLocation() {
-            return new Rectangle(
-                    location.getX(), location.getY(), location.getWidth(), location.getHeight());
-        }
-    }
 }
diff --git a/api/src/main/java/ai/djl/modality/cv/translator/YoloV8Translator.java b/api/src/main/java/ai/djl/modality/cv/translator/YoloV8Translator.java
index 7fb1fda62..7895ad15c 100644
--- a/api/src/main/java/ai/djl/modality/cv/translator/YoloV8Translator.java
+++ b/api/src/main/java/ai/djl/modality/cv/translator/YoloV8Translator.java
@@ -79,7 +79,10 @@ public class YoloV8Translator extends YoloV5Translator {
                     "Expected classes: " + (nClasses - 4) + ", got " + classes.size());
         }
 
-        ArrayList<IntermediateResult> intermediateResults = new ArrayList<>();
+        ArrayList<Rectangle> boxes = new ArrayList<>();
+        ArrayList<Float> scores = new ArrayList<>();
+        ArrayList<Integer> classIds = new ArrayList<>();
+
         // reverse order search in heap; searches through #maxBoxes for optimization when set
         for (int i = numberRows - 1; i > numberRows - maxBoxes; --i) {
             int index = i * nClasses;
@@ -101,12 +104,13 @@ public class YoloV8Translator extends YoloV5Translator {
                 float h = buf[index + 3];
                 Rectangle rect =
                         new Rectangle(Math.max(0, xPos - w / 2), Math.max(0, yPos - h / 2), w, h);
-                intermediateResults.add(
-                        new IntermediateResult(
-                                classes.get(maxIndex), maxClassProb, maxIndex, rect));
+                boxes.add(rect);
+                scores.add(maxClassProb);
+                classIds.add(maxIndex);
             }
         }
-        return nms(intermediateResults);
+
+        return nms(boxes, classIds, scores);
     }
 
     /** The builder for {@link YoloV8Translator}. */
diff --git a/api/src/main/java/ai/djl/repository/zoo/BaseModelLoader.java b/api/src/main/java/ai/djl/repository/zoo/BaseModelLoader.java
index 89de4c254..f2bdf29a7 100644
--- a/api/src/main/java/ai/djl/repository/zoo/BaseModelLoader.java
+++ b/api/src/main/java/ai/djl/repository/zoo/BaseModelLoader.java
@@ -195,11 +195,19 @@ public class BaseModelLoader implements ModelLoader {
             throws IOException {
         Model model = Model.newInstance(name, device, engine);
         if (block == null) {
-            String className = (String) arguments.get("blockFactory");
-            BlockFactory factory =
-                    ClassLoaderUtils.findImplementation(modelPath, BlockFactory.class, className);
-            if (factory != null) {
-                block = factory.newBlock(model, modelPath, arguments);
+            Object bf = arguments.get("blockFactory");
+            if (bf instanceof BlockFactory) {
+                block = ((BlockFactory) bf).newBlock(model, modelPath, arguments);
+            } else {
+                String className = (String) bf;
+                BlockFactory factory =
+                        ClassLoaderUtils.findImplementation(
+                                modelPath, BlockFactory.class, className);
+                if (factory != null) {
+                    block = factory.newBlock(model, modelPath, arguments);
+                } else if (className != null) {
+                    throw new IllegalArgumentException("Failed to load BlockFactory: " + className);
+                }
             }
         }
         if (block != null) {
diff --git a/api/src/main/java/ai/djl/util/Platform.java b/api/src/main/java/ai/djl/util/Platform.java
index 2ab25c2f3..dd31a3833 100644
--- a/api/src/main/java/ai/djl/util/Platform.java
+++ b/api/src/main/java/ai/djl/util/Platform.java
@@ -20,7 +20,9 @@ import org.slf4j.LoggerFactory;
 import java.io.IOException;
 import java.io.InputStream;
 import java.net.URL;
+import java.util.ArrayList;
 import java.util.Enumeration;
+import java.util.List;
 import java.util.Properties;
 
 /**
@@ -71,32 +73,50 @@ public final class Platform {
         }
 
         Platform systemPlatform = Platform.fromSystem(engine);
-        Platform placeholder = null;
+        List<Platform> availablePlatforms = new ArrayList<>();
         while (urls.hasMoreElements()) {
             URL url = urls.nextElement();
             Platform platform = Platform.fromUrl(url);
             platform.apiVersion = systemPlatform.apiVersion;
             if (platform.isPlaceholder()) {
-                placeholder = platform;
+                availablePlatforms.add(platform);
             } else if (platform.matches(systemPlatform)) {
                 logger.info("Found matching platform from: {}", url);
-                return platform;
+                availablePlatforms.add(platform);
             } else {
                 logger.info("Ignore mismatching platform from: {}", url);
             }
         }
-        if (placeholder != null) {
-            logger.info("Found placeholder platform from: {}", placeholder);
-            return placeholder;
-        }
-
-        if (systemPlatform.version == null) {
-            throw new AssertionError("No " + engine + " version found in property file.");
-        }
-        if (systemPlatform.apiVersion == null) {
-            throw new AssertionError("No " + engine + " djl_version found in property file.");
+        if (availablePlatforms.isEmpty()) {
+            if (systemPlatform.version == null) {
+                throw new AssertionError("No " + engine + " version found in property file.");
+            }
+            if (systemPlatform.apiVersion == null) {
+                throw new AssertionError("No " + engine + " djl_version found in property file.");
+            }
+            return systemPlatform;
+        } else if (availablePlatforms.size() == 1) {
+            Platform ret = availablePlatforms.get(0);
+            if (ret.isPlaceholder()) {
+                logger.info("Found placeholder platform from: {}", ret);
+            }
+            return ret;
         }
-        return systemPlatform;
+        availablePlatforms.sort(
+                (o1, o2) -> {
+                    if (o1.isPlaceholder()) {
+                        return 1;
+                    } else if (o2.isPlaceholder()) {
+                        return -1;
+                    }
+                    // cu121-precx11 > cu121 > cu118-precss11 > cpu-precxx11 > cpu
+                    int ret = o2.getFlavor().compareTo(o1.getFlavor());
+                    if (ret == 0) {
+                        return o2.getVersion().compareTo(o1.getVersion());
+                    }
+                    return ret;
+                });
+        return availablePlatforms.get(0);
     }
 
     /**
diff --git a/api/src/main/java/ai/djl/util/cuda/CudaUtils.java b/api/src/main/java/ai/djl/util/cuda/CudaUtils.java
index 859596a6b..b234eb479 100644
--- a/api/src/main/java/ai/djl/util/cuda/CudaUtils.java
+++ b/api/src/main/java/ai/djl/util/cuda/CudaUtils.java
@@ -139,11 +139,13 @@ public final class CudaUtils {
      */
     public static String getComputeCapability(int device) {
         if (Boolean.getBoolean("ai.djl.util.cuda.fork")) {
-            String[] ret = execute(device);
-            if (ret.length != 3) {
-                throw new IllegalArgumentException(ret[0]);
+            if (gpuInfo == null) { // NOPMD
+                gpuInfo = execute(-1);
             }
-            return ret[0];
+            if (device >= gpuInfo.length - 2) {
+                throw new IllegalArgumentException("Invalid device: " + device);
+            }
+            return gpuInfo[device + 2];
         }
 
         if (LIB == null) {
@@ -214,7 +216,12 @@ public final class CudaUtils {
                 return;
             }
             int cudaVersion = getCudaVersion();
-            System.out.println(gpuCount + "," + cudaVersion);
+            StringBuilder sb = new StringBuilder();
+            sb.append(gpuCount).append(',').append(cudaVersion);
+            for (int i = 0; i < gpuCount; ++i) {
+                sb.append(',').append(getComputeCapability(i));
+            }
+            System.out.println(sb);
             return;
         }
         try {
diff --git a/examples/src/main/java/ai/djl/examples/training/TrainTimeSeries.java b/examples/src/main/java/ai/djl/examples/training/TrainTimeSeries.java
index b9abb016e..66f734e9c 100644
--- a/examples/src/main/java/ai/djl/examples/training/TrainTimeSeries.java
+++ b/examples/src/main/java/ai/djl/examples/training/TrainTimeSeries.java
@@ -120,10 +120,10 @@ public final class TrainTimeSeries {
                 inputShapes[6] =
                         new Shape(
                                 1,
-                                predictionLength,
+                                contextLength,
                                 TimeFeature.timeFeaturesFromFreqStr(freq).size() + 1);
-                inputShapes[7] = new Shape(1, predictionLength);
-                inputShapes[8] = new Shape(1, predictionLength);
+                inputShapes[7] = new Shape(1, contextLength);
+                inputShapes[8] = new Shape(1, contextLength);
                 trainer.initialize(inputShapes);
 
                 EasyTrain.fit(trainer, arguments.getEpoch(), trainSet, null);
@@ -147,6 +147,7 @@ public final class TrainTimeSeries {
 
             Map<String, Object> arguments = new ConcurrentHashMap<>();
             arguments.put("prediction_length", predictionLength);
+            arguments.put("context_length", predictionNetwork.getContextLength());
             arguments.put("freq", freq);
             arguments.put("use_" + FieldName.FEAT_DYNAMIC_REAL.name().toLowerCase(), false);
             arguments.put("use_" + FieldName.FEAT_STATIC_CAT.name().toLowerCase(), true);
@@ -239,6 +240,7 @@ public final class TrainTimeSeries {
                         .setFreq(freq)
                         .setPredictionLength(predictionLength)
                         .optDistrOutput(distributionOutput)
+                        .optContextLength(8)
                         .optUseFeatStaticCat(true);
         return training ? builder.buildTrainingNetwork() : builder.buildPredictionNetwork();
     }
diff --git a/examples/src/main/java/ai/djl/examples/training/TrainWithOptimizers.java b/examples/src/main/java/ai/djl/examples/training/TrainWithOptimizers.java
index b9f254eea..d8b6c9dcc 100644
--- a/examples/src/main/java/ai/djl/examples/training/TrainWithOptimizers.java
+++ b/examples/src/main/java/ai/djl/examples/training/TrainWithOptimizers.java
@@ -15,7 +15,6 @@ package ai.djl.examples.training;
 import ai.djl.Model;
 import ai.djl.ModelException;
 import ai.djl.basicdataset.cv.classification.Cifar10;
-import ai.djl.basicmodelzoo.BasicModelZoo;
 import ai.djl.basicmodelzoo.cv.classification.ResNetV1;
 import ai.djl.examples.training.util.Arguments;
 import ai.djl.metric.Metrics;
@@ -25,10 +24,6 @@ import ai.djl.modality.cv.transform.Normalize;
 import ai.djl.modality.cv.transform.ToTensor;
 import ai.djl.ndarray.types.Shape;
 import ai.djl.nn.Block;
-import ai.djl.nn.Blocks;
-import ai.djl.nn.SequentialBlock;
-import ai.djl.nn.SymbolBlock;
-import ai.djl.nn.core.Linear;
 import ai.djl.repository.zoo.Criteria;
 import ai.djl.training.DefaultTrainingConfig;
 import ai.djl.training.EasyTrain;
@@ -53,7 +48,6 @@ import org.apache.commons.cli.ParseException;
 
 import java.io.IOException;
 import java.util.Arrays;
-import java.util.Map;
 
 /** This example features sample usage of a variety of optimizers to train Cifar10. */
 public final class TrainWithOptimizers {
@@ -97,50 +91,15 @@ public final class TrainWithOptimizers {
     }
 
     private static Model getModel(Arguments arguments) throws IOException, ModelException {
-        boolean isSymbolic = arguments.isSymbolic();
         boolean preTrained = arguments.isPreTrained();
-        Map<String, String> options = arguments.getCriteria();
         Criteria.Builder<Image, Classifications> builder =
                 Criteria.builder()
                         .setTypes(Image.class, Classifications.class)
                         .optEngine(arguments.getEngine())
-                        .optProgress(new ProgressBar())
-                        .optArtifactId("resnet");
-        if (isSymbolic) {
-            // currently only MxEngine support removeLastBlock
-            builder.optGroupId("ai.djl.mxnet");
-            if (options == null) {
-                builder.optFilter("layers", "50");
-                builder.optFilter("flavor", "v1");
-            } else {
-                builder.optFilters(options);
-            }
-
-            Model model = builder.build().loadModel();
-            SequentialBlock newBlock = new SequentialBlock();
-            SymbolBlock block = (SymbolBlock) model.getBlock();
-            block.removeLastBlock();
-            newBlock.add(block);
-            // the original model don't include the flatten
-            // so apply the flatten here
-            newBlock.add(Blocks.batchFlattenBlock());
-            newBlock.add(Linear.builder().setUnits(10).build());
-            model.setBlock(newBlock);
-            if (!preTrained) {
-                model.getBlock().clear();
-            }
-            return model;
-        }
+                        .optProgress(new ProgressBar());
         // imperative resnet50
         if (preTrained) {
-            builder.optGroupId(BasicModelZoo.GROUP_ID);
-            if (options == null) {
-                builder.optFilter("layers", "50");
-                builder.optFilter("flavor", "v1");
-                builder.optFilter("dataset", "cifar10");
-            } else {
-                builder.optFilters(options);
-            }
+            builder.optModelUrls("djl://ai.djl.zoo/resnet/0.0.2/resnetv1");
             // load pre-trained imperative ResNet50 from DJL model zoo
             return builder.build().loadModel();
         } else {
diff --git a/examples/src/main/java/ai/djl/examples/training/transferlearning/TrainResnetWithCifar10.java b/examples/src/main/java/ai/djl/examples/training/transferlearning/TrainResnetWithCifar10.java
index e6e7e33d9..00317f82b 100644
--- a/examples/src/main/java/ai/djl/examples/training/transferlearning/TrainResnetWithCifar10.java
+++ b/examples/src/main/java/ai/djl/examples/training/transferlearning/TrainResnetWithCifar10.java
@@ -12,11 +12,9 @@
  */
 package ai.djl.examples.training.transferlearning;
 
-import ai.djl.Application;
 import ai.djl.Model;
 import ai.djl.ModelException;
 import ai.djl.basicdataset.cv.classification.Cifar10;
-import ai.djl.basicmodelzoo.BasicModelZoo;
 import ai.djl.basicmodelzoo.cv.classification.ResNetV1;
 import ai.djl.examples.training.util.Arguments;
 import ai.djl.inference.Predictor;
@@ -30,10 +28,7 @@ import ai.djl.modality.cv.translator.ImageClassificationTranslator;
 import ai.djl.ndarray.NDManager;
 import ai.djl.ndarray.types.Shape;
 import ai.djl.nn.Block;
-import ai.djl.nn.Blocks;
-import ai.djl.nn.SequentialBlock;
-import ai.djl.nn.SymbolBlock;
-import ai.djl.nn.core.Linear;
+import ai.djl.nn.BlockFactory;
 import ai.djl.repository.zoo.Criteria;
 import ai.djl.repository.zoo.ZooModel;
 import ai.djl.training.DefaultTrainingConfig;
@@ -55,7 +50,6 @@ import org.slf4j.LoggerFactory;
 import java.io.IOException;
 import java.nio.file.Path;
 import java.nio.file.Paths;
-import java.util.Map;
 
 /**
  * An example of training an image classification (ResNet for Cifar10) model.
@@ -112,8 +106,7 @@ public final class TrainResnetWithCifar10 {
                 Path modelPath = Paths.get("build/model");
                 model.save(modelPath, "resnetv1");
 
-                Classifications classifications =
-                        testSaveParameters(model.getBlock(), modelPath, arguments);
+                Classifications classifications = testSaveParameters(modelPath, arguments);
                 logger.info("Predict result: {}", classifications.topK(3));
                 return result;
             }
@@ -121,50 +114,15 @@ public final class TrainResnetWithCifar10 {
     }
 
     private static Model getModel(Arguments arguments) throws IOException, ModelException {
-        boolean isSymbolic = arguments.isSymbolic();
         boolean preTrained = arguments.isPreTrained();
-        Map<String, String> options = arguments.getCriteria();
         Criteria.Builder<Image, Classifications> builder =
                 Criteria.builder()
-                        .optApplication(Application.CV.IMAGE_CLASSIFICATION)
                         .setTypes(Image.class, Classifications.class)
                         .optEngine(arguments.getEngine())
-                        .optProgress(new ProgressBar())
-                        .optArtifactId("resnet");
-        if (isSymbolic) {
-            // load the model
-            builder.optGroupId("ai.djl.mxnet");
-            if (options == null) {
-                builder.optFilter("layers", "50");
-                builder.optFilter("flavor", "v1");
-            } else {
-                builder.optFilters(options);
-            }
-            Model model = builder.build().loadModel();
-            SequentialBlock newBlock = new SequentialBlock();
-            SymbolBlock block = (SymbolBlock) model.getBlock();
-            block.removeLastBlock();
-            newBlock.add(block);
-            // the original model don't include the flatten
-            // so apply the flatten here
-            newBlock.add(Blocks.batchFlattenBlock());
-            newBlock.add(Linear.builder().setUnits(10).build());
-            model.setBlock(newBlock);
-            if (!preTrained) {
-                model.getBlock().clear();
-            }
-            return model;
-        }
+                        .optProgress(new ProgressBar());
         // imperative resnet50
         if (preTrained) {
-            builder.optGroupId(BasicModelZoo.GROUP_ID);
-            if (options == null) {
-                builder.optFilter("layers", "50");
-                builder.optFilter("flavor", "v1");
-                builder.optFilter("dataset", "cifar10");
-            } else {
-                builder.optFilters(options);
-            }
+            builder.optModelUrls("djl://ai.djl.zoo/resnet/0.0.2/resnetv1");
             // load pre-trained imperative ResNet50 from DJL model zoo
             return builder.build().loadModel();
         } else {
@@ -181,7 +139,7 @@ public final class TrainResnetWithCifar10 {
         }
     }
 
-    private static Classifications testSaveParameters(Block block, Path path, Arguments arguments)
+    static Classifications testSaveParameters(Path path, Arguments arguments)
             throws IOException, ModelException, TranslateException {
         String synsetUrl =
                 "https://mlrepo.djl.ai/model/cv/image_classification/ai/djl/mxnet/synset_cifar10.txt";
@@ -192,6 +150,13 @@ public final class TrainResnetWithCifar10 {
                         .optSynsetUrl(synsetUrl)
                         .optApplySoftmax(true)
                         .build();
+        BlockFactory resnetFactory =
+                (model, modelPath, arguments1) ->
+                        ResNetV1.builder()
+                                .setImageShape(new Shape(3, 32, 32))
+                                .setNumLayers(50)
+                                .setOutSize(10)
+                                .build();
 
         Image img = ImageFactory.getInstance().fromUrl("src/test/resources/airplane1.png");
 
@@ -201,7 +166,7 @@ public final class TrainResnetWithCifar10 {
                         .optModelPath(path)
                         .optEngine(arguments.getEngine())
                         .optTranslator(translator)
-                        .optBlock(block)
+                        .optArgument("blockFactory", resnetFactory)
                         .optModelName("resnetv1")
                         .build();
 
diff --git a/examples/src/main/java/ai/djl/examples/training/util/Arguments.java b/examples/src/main/java/ai/djl/examples/training/util/Arguments.java
index 5a0048226..673698df1 100644
--- a/examples/src/main/java/ai/djl/examples/training/util/Arguments.java
+++ b/examples/src/main/java/ai/djl/examples/training/util/Arguments.java
@@ -33,7 +33,6 @@ public class Arguments {
     protected int epoch;
     protected int batchSize;
     protected int maxGpus;
-    protected boolean isSymbolic;
     protected boolean preTrained;
     protected String outputDir;
     protected long limit;
@@ -60,7 +59,6 @@ public class Arguments {
         } else {
             batchSize = maxGpus > 0 ? 32 * maxGpus : 32;
         }
-        isSymbolic = cmd.hasOption("symbolic-model");
         preTrained = cmd.hasOption("pre-trained");
 
         if (cmd.hasOption("output-dir")) {
@@ -126,12 +124,6 @@ public class Arguments {
                         .argName("MAXGPUS")
                         .desc("Max number of GPUs to use for training")
                         .build());
-        options.addOption(
-                Option.builder("s")
-                        .longOpt("symbolic-model")
-                        .argName("SYMBOLIC")
-                        .desc("Use symbolic model, use imperative model if false")
-                        .build());
         options.addOption(
                 Option.builder("p")
                         .longOpt("pre-trained")
@@ -190,10 +182,6 @@ public class Arguments {
         return Engine.getEngine(engine).getDevices(maxGpus);
     }
 
-    public boolean isSymbolic() {
-        return isSymbolic;
-    }
-
     public boolean isPreTrained() {
         return preTrained;
     }
diff --git a/extensions/opencv/src/main/java/ai/djl/opencv/OpenCVImage.java b/extensions/opencv/src/main/java/ai/djl/opencv/OpenCVImage.java
index b6a229271..1cbacd634 100644
--- a/extensions/opencv/src/main/java/ai/djl/opencv/OpenCVImage.java
+++ b/extensions/opencv/src/main/java/ai/djl/opencv/OpenCVImage.java
@@ -356,27 +356,37 @@ class OpenCVImage implements Image {
         float b = RandomUtils.nextFloat();
         int imageWidth = img.getWidth();
         int imageHeight = img.getHeight();
-        int x = (int) (mask.getX() * imageWidth);
-        int y = (int) (mask.getY() * imageHeight);
-        float[][] probDist = mask.getProbDist();
-        // Correct some coordinates of box when going out of image
-        if (x < 0) {
-            x = 0;
-        }
-        if (y < 0) {
-            y = 0;
+        int x = 0;
+        int y = 0;
+        int w = imageWidth;
+        int h = imageHeight;
+        if (!mask.isFullImageMask()) {
+            x = (int) (mask.getX() * imageWidth);
+            y = (int) (mask.getY() * imageHeight);
+            w = (int) (mask.getWidth() * imageWidth);
+            h = (int) (mask.getHeight() * imageHeight);
+            // Correct some coordinates of box when going out of image
+            if (x < 0) {
+                x = 0;
+            }
+            if (y < 0) {
+                y = 0;
+            }
         }
+        float[][] probDist = mask.getProbDist();
 
         BufferedImage maskImage =
-                new BufferedImage(probDist.length, probDist[0].length, BufferedImage.TYPE_INT_ARGB);
-        for (int xCor = 0; xCor < probDist.length; xCor++) {
-            for (int yCor = 0; yCor < probDist[xCor].length; yCor++) {
-                float opacity = probDist[xCor][yCor] * 0.8f;
-                maskImage.setRGB(xCor, yCor, new Color(r, g, b, opacity).getRGB());
+                new BufferedImage(probDist[0].length, probDist.length, BufferedImage.TYPE_INT_ARGB);
+        for (int yCor = 0; yCor < probDist.length; yCor++) {
+            for (int xCor = 0; xCor < probDist[0].length; xCor++) {
+                float opacity = probDist[yCor][xCor] * 0.8f;
+                maskImage.setRGB(xCor, yCor, new Color(r, g, b, opacity).darker().getRGB());
             }
         }
+        java.awt.Image scaled = maskImage.getScaledInstance(w, h, java.awt.Image.SCALE_SMOOTH);
+
         Graphics2D gR = (Graphics2D) img.getGraphics();
-        gR.drawImage(maskImage, x, y, null);
+        gR.drawImage(scaled, x, y, null);
         gR.dispose();
     }
 
diff --git a/extensions/timeseries/src/main/java/ai/djl/timeseries/transform/feature/Feature.java b/extensions/timeseries/src/main/java/ai/djl/timeseries/transform/feature/Feature.java
index 0ade03d3f..f7e1f6255 100644
--- a/extensions/timeseries/src/main/java/ai/djl/timeseries/transform/feature/Feature.java
+++ b/extensions/timeseries/src/main/java/ai/djl/timeseries/transform/feature/Feature.java
@@ -119,14 +119,14 @@ public final class Feature {
         if (!freq.matches("\\d+.*")) {
             sb.insert(0, 1);
         }
-        sb.insert(0, "P");
-        String formattedFreq = sb.toString();
 
         TemporalAmount timeFreq;
         if (freq.endsWith("H") || freq.endsWith("T") || freq.endsWith("S")) {
-            timeFreq = Duration.parse(formattedFreq);
+            sb.insert(0, "PT");
+            timeFreq = Duration.parse(sb.toString());
         } else {
-            timeFreq = Period.parse(formattedFreq);
+            sb.insert(0, "P");
+            timeFreq = Period.parse(sb.toString());
         }
 
         List<LocalDateTime> index = new ArrayList<>();
diff --git a/extensions/timeseries/src/main/java/ai/djl/timeseries/translator/BaseTimeSeriesTranslator.java b/extensions/timeseries/src/main/java/ai/djl/timeseries/translator/BaseTimeSeriesTranslator.java
index c84e874a3..d34bcfb6f 100644
--- a/extensions/timeseries/src/main/java/ai/djl/timeseries/translator/BaseTimeSeriesTranslator.java
+++ b/extensions/timeseries/src/main/java/ai/djl/timeseries/translator/BaseTimeSeriesTranslator.java
@@ -39,8 +39,7 @@ public abstract class BaseTimeSeriesTranslator implements Translator<TimeSeriesD
         this.batchifier = builder.batchifier;
         this.freq = builder.freq;
         this.predictionLength = builder.predictionLength;
-        // TODO: for inferring
-        this.contextLength = builder.predictionLength;
+        this.contextLength = builder.contextLength;
     }
 
     /** {@inheritDoc} */
@@ -57,6 +56,7 @@ public abstract class BaseTimeSeriesTranslator implements Translator<TimeSeriesD
     public abstract static class BaseBuilder<T extends BaseBuilder<T>> {
         protected Batchifier batchifier = Batchifier.STACK;
         protected int predictionLength;
+        protected int contextLength;
 
         protected String freq;
 
@@ -82,6 +82,8 @@ public abstract class BaseTimeSeriesTranslator implements Translator<TimeSeriesD
                 throw new IllegalArgumentException(
                         "The value of `prediction_length` should be > 0");
             }
+            this.contextLength =
+                    ArgumentsUtil.intValue(arguments, "context_length", predictionLength);
             if (arguments.containsKey("batchifier")) {
                 batchifier = Batchifier.fromString((String) arguments.get("batchifier"));
             }
diff --git a/extensions/tokenizers/rust/src/layers/cublaslt.rs b/extensions/tokenizers/rust/src/layers/cublaslt.rs
index 0ecca7000..4fe5431e4 100644
--- a/extensions/tokenizers/rust/src/layers/cublaslt.rs
+++ b/extensions/tokenizers/rust/src/layers/cublaslt.rs
@@ -2,8 +2,8 @@
 pub use cudarc::cublaslt::Activation;
 
 use crate::layers::HiddenAct;
-use candle::{CpuStorage, Device, Layout, Result, Shape, Storage, Tensor};
 use candle::backend::BackendStorage;
+use candle::{CpuStorage, Device, Layout, Result, Shape, Storage, Tensor};
 
 #[cfg(feature = "cuda")]
 use candle::cuda_backend::WrapErr;
diff --git a/extensions/tokenizers/rust/src/layers/linear.rs b/extensions/tokenizers/rust/src/layers/linear.rs
index 966a672a6..2b67d777b 100644
--- a/extensions/tokenizers/rust/src/layers/linear.rs
+++ b/extensions/tokenizers/rust/src/layers/linear.rs
@@ -20,7 +20,12 @@ pub struct Linear {
 }
 
 impl Linear {
-    pub fn load(vb: VarBuilder, in_dim: usize, out_dim: usize, act: Option<HiddenAct>) -> Result<Self> {
+    pub fn load(
+        vb: VarBuilder,
+        in_dim: usize,
+        out_dim: usize,
+        act: Option<HiddenAct>,
+    ) -> Result<Self> {
         Ok(Self {
             weight: vb.get((out_dim, in_dim), "weight")?,
             bias: Some(vb.get(out_dim, "bias")?),
diff --git a/extensions/tokenizers/rust/src/layers/mod.rs b/extensions/tokenizers/rust/src/layers/mod.rs
index 378588fdf..3f0cd444c 100644
--- a/extensions/tokenizers/rust/src/layers/mod.rs
+++ b/extensions/tokenizers/rust/src/layers/mod.rs
@@ -2,5 +2,4 @@
 mod cublaslt;
 mod linear;
 
-pub use cublaslt::get_cublas_lt_wrapper;
 pub use linear::{HiddenAct, Linear};
diff --git a/extensions/tokenizers/rust/src/lib.rs b/extensions/tokenizers/rust/src/lib.rs
index 5e37294a8..2d7a93e23 100644
--- a/extensions/tokenizers/rust/src/lib.rs
+++ b/extensions/tokenizers/rust/src/lib.rs
@@ -16,8 +16,8 @@ mod ndarray;
 
 #[cfg(feature = "cuda")]
 mod compute_cap;
-mod models;
 mod layers;
+mod models;
 
 extern crate tokenizers as tk;
 
@@ -30,7 +30,6 @@ use jni::errors::Error;
 use jni::objects::{
     JClass, JLongArray, JMethodID, JObject, JObjectArray, JString, JValue, ReleaseMode,
 };
-#[cfg(feature = "cuda")]
 use jni::sys::JNI_FALSE;
 
 use jni::sys::{jboolean, jint, jlong, jobjectArray, jsize, jvalue, JNI_TRUE};
@@ -811,7 +810,7 @@ pub extern "system" fn Java_ai_djl_engine_rust_RustLibrary_isCudaAvailable<'loca
     }
     #[cfg(not(feature = "cuda"))]
     {
-        JNI_TRUE
+        JNI_FALSE
     }
 }
 
diff --git a/extensions/tokenizers/rust/src/models/bert.rs b/extensions/tokenizers/rust/src/models/bert.rs
index 7d53e96fb..c241232d3 100644
--- a/extensions/tokenizers/rust/src/models/bert.rs
+++ b/extensions/tokenizers/rust/src/models/bert.rs
@@ -1,5 +1,5 @@
-use crate::models::Model;
 use crate::layers::Linear;
+use crate::models::Model;
 use candle::{Device, Result, Tensor};
 use candle_nn::{embedding, Embedding, Module, VarBuilder};
 use candle_transformers::models::with_tracing::{layer_norm, LayerNorm};
@@ -359,7 +359,12 @@ struct BertIntermediate {
 
 impl BertIntermediate {
     fn load(vb: VarBuilder, config: &BertConfig) -> Result<Self> {
-        let dense = Linear::load(vb.pp("dense"), config.hidden_size, config.intermediate_size, None)?;
+        let dense = Linear::load(
+            vb.pp("dense"),
+            config.hidden_size,
+            config.intermediate_size,
+            None,
+        )?;
         Ok(Self {
             dense,
             intermediate_act: HiddenActLayer::new(config.hidden_act),
@@ -387,7 +392,12 @@ struct BertOutput {
 
 impl BertOutput {
     fn load(vb: VarBuilder, config: &BertConfig) -> Result<Self> {
-        let dense = Linear::load(vb.pp("dense"), config.intermediate_size, config.hidden_size, None)?;
+        let dense = Linear::load(
+            vb.pp("dense"),
+            config.intermediate_size,
+            config.hidden_size,
+            None,
+        )?;
         let layer_norm = layer_norm(
             config.hidden_size,
             config.layer_norm_eps,
diff --git a/extensions/tokenizers/rust/src/models/mod.rs b/extensions/tokenizers/rust/src/models/mod.rs
index b82f00670..a0b52daa4 100644
--- a/extensions/tokenizers/rust/src/models/mod.rs
+++ b/extensions/tokenizers/rust/src/models/mod.rs
@@ -4,7 +4,7 @@ mod distilbert;
 use crate::ndarray::as_data_type;
 use crate::{cast_handle, drop_handle, to_handle, to_string_array};
 use bert::{BertConfig, BertModel};
-use candle::{Device, DType, Result, Tensor};
+use candle::{DType, Device, Result, Tensor};
 use candle_nn::VarBuilder;
 use distilbert::{DistilBertConfig, DistilBertModel};
 use jni::objects::{JLongArray, JObject, JString, ReleaseMode};


Test patch:
diff --git a/api/src/test/java/ai/djl/util/PlatformTest.java b/api/src/test/java/ai/djl/util/PlatformTest.java
index 039cbbea8..3e20609a0 100644
--- a/api/src/test/java/ai/djl/util/PlatformTest.java
+++ b/api/src/test/java/ai/djl/util/PlatformTest.java
@@ -12,16 +12,23 @@
  */
 package ai.djl.util;
 
+import ai.djl.util.cuda.CudaUtils;
+
 import org.testng.Assert;
 import org.testng.annotations.AfterClass;
 import org.testng.annotations.Test;
 
 import java.io.BufferedWriter;
 import java.io.IOException;
+import java.lang.reflect.Field;
 import java.net.URL;
+import java.net.URLClassLoader;
+import java.nio.charset.StandardCharsets;
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.nio.file.Paths;
+import java.util.jar.JarEntry;
+import java.util.jar.JarOutputStream;
 
 public class PlatformTest {
 
@@ -94,6 +101,43 @@ public class PlatformTest {
         Assert.assertFalse(platform.matches(system));
     }
 
+    @Test
+    public void testDetectPlatform() throws IOException, ReflectiveOperationException {
+        Path dir = Paths.get("build/tmp/");
+        Files.createDirectories(dir);
+        Platform system = Platform.fromSystem();
+        String classifier = system.getClassifier();
+        createZipFile(0, "1.0", "cpu", classifier, true);
+        createZipFile(1, "1.0", "cpu", classifier, false);
+        createZipFile(2, "1.0", "cpu-precxx11", classifier, false);
+        createZipFile(3, "1.0", "cu117", classifier, false);
+        createZipFile(4, "1.0", "cu117-precxx11", classifier, false);
+        createZipFile(5, "1.0", "cu999", classifier, false);
+        createZipFile(6, "1.0", "cu999-precxx11", classifier, false);
+        createZipFile(7, "99.99", "cu999-precxx11", classifier, false);
+        System.setProperty("ai.djl.util.cuda.fork", "true");
+        try {
+            String[] gpuInfo = new String[] {"1", "99990", "90"};
+            Field field = CudaUtils.class.getDeclaredField("gpuInfo");
+            field.setAccessible(true);
+            field.set(null, gpuInfo);
+            URL[] urls = new URL[8];
+            for (int i = 0; i < 8; ++i) {
+                urls[i] = dir.resolve(i + ".jar").toUri().toURL();
+            }
+            URLClassLoader cl = new URLClassLoader(urls);
+            Thread.currentThread().setContextClassLoader(cl);
+
+            Platform detected = Platform.detectPlatform("pytorch");
+            Assert.assertEquals(detected.getFlavor(), "cu999-precxx11");
+
+            field.set(null, null);
+        } finally {
+            System.clearProperty("ai.djl.util.cuda.fork");
+            Thread.currentThread().setContextClassLoader(null);
+        }
+    }
+
     private URL createPropertyFile(String content) throws IOException {
         Path dir = Paths.get("build/tmp/testFile/");
         Files.createDirectories(dir);
@@ -104,4 +148,24 @@ public class PlatformTest {
         }
         return file.toUri().toURL();
     }
+
+    private void createZipFile(
+            int index, String version, String flavor, String classifier, boolean placeHolder)
+            throws IOException {
+        Path file = Paths.get("build/tmp/" + index + ".jar");
+        try (JarOutputStream jos = new JarOutputStream(Files.newOutputStream(file))) {
+            JarEntry entry = new JarEntry("native/lib/pytorch.properties");
+            jos.putNextEntry(entry);
+            if (placeHolder) {
+                jos.write("placeholder=true\nversion=2.3.1".getBytes(StandardCharsets.UTF_8));
+            } else {
+                jos.write("version=".getBytes(StandardCharsets.UTF_8));
+                jos.write(version.getBytes(StandardCharsets.UTF_8));
+                jos.write("\nflavor=".getBytes(StandardCharsets.UTF_8));
+                jos.write(flavor.getBytes(StandardCharsets.UTF_8));
+                jos.write("\nclassifier=".getBytes(StandardCharsets.UTF_8));
+                jos.write(classifier.getBytes(StandardCharsets.UTF_8));
+            }
+        }
+    }
 }
diff --git a/api/src/test/java/ai/djl/util/SecurityManagerTest.java b/api/src/test/java/ai/djl/util/SecurityManagerTest.java
index 1e9eb17f6..fbe09975d 100644
--- a/api/src/test/java/ai/djl/util/SecurityManagerTest.java
+++ b/api/src/test/java/ai/djl/util/SecurityManagerTest.java
@@ -12,14 +12,11 @@
  */
 package ai.djl.util;
 
-import ai.djl.util.cuda.CudaUtils;
-
 import org.testng.Assert;
 import org.testng.annotations.AfterTest;
 import org.testng.annotations.BeforeTest;
 import org.testng.annotations.Test;
 
-import java.io.FilePermission;
 import java.security.Permission;
 
 public class SecurityManagerTest {
@@ -59,26 +56,4 @@ public class SecurityManagerTest {
         Assert.assertEquals(Utils.getenv("TEST", "test"), "test");
         Assert.assertEquals(Utils.getenv().size(), 0);
     }
-
-    @Test
-    public void testCudaUtils() {
-        // Disable access to the cudart library files
-        SecurityManager sm =
-                new SecurityManager() {
-                    @Override
-                    public void checkPermission(Permission perm) {
-                        if (perm instanceof FilePermission && perm.getName().contains("cudart")) {
-                            throw new SecurityException(
-                                    "Don't have permission to read file: " + perm.getName());
-                        }
-                    }
-                };
-        System.setSecurityManager(sm);
-        try {
-            Assert.assertFalse(CudaUtils.hasCuda());
-            Assert.assertEquals(CudaUtils.getGpuCount(), 0);
-        } finally {
-            System.setSecurityManager(null);
-        }
-    }
 }
diff --git a/api/src/test/java/ai/djl/util/cuda/CudaUtilsTest.java b/api/src/test/java/ai/djl/util/cuda/CudaUtilsTest.java
index df135938a..f6ecc66ac 100644
--- a/api/src/test/java/ai/djl/util/cuda/CudaUtilsTest.java
+++ b/api/src/test/java/ai/djl/util/cuda/CudaUtilsTest.java
@@ -55,6 +55,8 @@ public class CudaUtilsTest {
         System.setProperty("ai.djl.util.cuda.fork", "true");
         try {
             testCudaUtils();
+            CudaUtils.main(new String[0]);
+            CudaUtils.main(new String[] {"-1"});
         } finally {
             System.clearProperty("ai.djl.util.cuda.fork");
         }
diff --git a/examples/src/test/java/ai/djl/examples/inference/MaskDetectionTest.java b/examples/src/test/java/ai/djl/examples/inference/MaskDetectionTest.java
index 551920bff..02ef7f6f8 100644
--- a/examples/src/test/java/ai/djl/examples/inference/MaskDetectionTest.java
+++ b/examples/src/test/java/ai/djl/examples/inference/MaskDetectionTest.java
@@ -35,6 +35,7 @@ public class MaskDetectionTest {
     @Test
     public void testMaskDetection() throws ModelException, TranslateException, IOException {
         TestRequirements.linux();
+        TestRequirements.notGpu();
 
         DetectedObjects result = MaskDetection.predict();
         logger.info("{}", result);
diff --git a/examples/src/test/java/ai/djl/examples/inference/ObjectDetectionWithTensorflowSavedModelTest.java b/examples/src/test/java/ai/djl/examples/inference/ObjectDetectionWithTensorflowSavedModelTest.java
index afb300551..cb8448113 100644
--- a/examples/src/test/java/ai/djl/examples/inference/ObjectDetectionWithTensorflowSavedModelTest.java
+++ b/examples/src/test/java/ai/djl/examples/inference/ObjectDetectionWithTensorflowSavedModelTest.java
@@ -35,6 +35,7 @@ public class ObjectDetectionWithTensorflowSavedModelTest {
         // throttling and will fail the test.
         TestRequirements.linux();
         TestRequirements.nightly();
+        TestRequirements.notGpu();
 
         DetectedObjects result = ObjectDetectionWithTensorflowSavedModel.predict();
 
diff --git a/examples/src/test/java/ai/djl/examples/inference/SuperResolutionTest.java b/examples/src/test/java/ai/djl/examples/inference/SuperResolutionTest.java
index 2c1a363a0..e40c2e6ad 100644
--- a/examples/src/test/java/ai/djl/examples/inference/SuperResolutionTest.java
+++ b/examples/src/test/java/ai/djl/examples/inference/SuperResolutionTest.java
@@ -32,6 +32,7 @@ public class SuperResolutionTest {
     @Test
     public void testSuperResolution() throws ModelException, TranslateException, IOException {
         TestRequirements.linux();
+        TestRequirements.notGpu();
 
         String imagePath = "src/test/resources/";
         Image fox = ImageFactory.getInstance().fromFile(Paths.get(imagePath + "fox.png"));
diff --git a/examples/src/test/java/ai/djl/examples/inference/UniversalSentenceEncoderTest.java b/examples/src/test/java/ai/djl/examples/inference/UniversalSentenceEncoderTest.java
index 9784b47b3..a85564fb6 100644
--- a/examples/src/test/java/ai/djl/examples/inference/UniversalSentenceEncoderTest.java
+++ b/examples/src/test/java/ai/djl/examples/inference/UniversalSentenceEncoderTest.java
@@ -29,6 +29,7 @@ public class UniversalSentenceEncoderTest {
     public void testSentimentAnalysis() throws ModelException, TranslateException, IOException {
         TestRequirements.linux();
         TestRequirements.nightly();
+        TestRequirements.notGpu();
 
         List<String> inputs = new ArrayList<>();
         inputs.add("The quick brown fox jumps over the lazy dog.");
diff --git a/examples/src/test/java/ai/djl/examples/inference/Yolov8DetectionTest.java b/examples/src/test/java/ai/djl/examples/inference/Yolov8DetectionTest.java
index 9fb9cc079..de46c8967 100644
--- a/examples/src/test/java/ai/djl/examples/inference/Yolov8DetectionTest.java
+++ b/examples/src/test/java/ai/djl/examples/inference/Yolov8DetectionTest.java
@@ -28,7 +28,7 @@ public class Yolov8DetectionTest {
 
     @Test
     public void testYolov8Detection() throws ModelException, TranslateException, IOException {
-        TestRequirements.linux();
+        TestRequirements.notGpu();
 
         DetectedObjects result = Yolov8Detection.predict();
 
diff --git a/examples/src/test/java/ai/djl/examples/training/TrainResNetTest.java b/examples/src/test/java/ai/djl/examples/training/TrainResNetTest.java
index f33be26de..cde2e1752 100644
--- a/examples/src/test/java/ai/djl/examples/training/TrainResNetTest.java
+++ b/examples/src/test/java/ai/djl/examples/training/TrainResNetTest.java
@@ -26,40 +26,21 @@ import org.testng.annotations.Test;
 import java.io.IOException;
 
 public class TrainResNetTest {
+
     private static final int SEED = 1234;
 
     @Test
     public void testTrainResNet() throws ModelException, IOException, TranslateException {
-        TestRequirements.linux();
-
-        // Limit max 4 gpu for cifar10 training to make it converge faster.
-        // and only train 10 batch for unit test.
-        // only MXNet support symbolic model
-        String[] args = {"-e", "2", "-g", "4", "-m", "10", "-s", "-p", "--engine", "MXNet"};
-
-        TrainingResult result = TrainResnetWithCifar10.runExample(args);
-        Assert.assertNotNull(result);
-    }
-
-    @Test
-    public void testTrainResNetSymbolicNightly()
-            throws ModelException, IOException, TranslateException {
-        TestRequirements.linux();
         TestRequirements.nightly();
-        TestRequirements.gpu("MXNet");
+        TestRequirements.gpu("PyTorch", 1);
 
         // Limit max 4 gpu for cifar10 training to make it converge faster.
         // and only train 10 batch for unit test.
         // only MXNet support symbolic model
-        String[] args = {"-e", "10", "-g", "4", "-s", "-p", "--engine", "MXNet"};
-
-        Engine.getEngine("MXNet").setRandomSeed(SEED);
+        String[] args = {"-e", "2", "-g", "4", "-m", "10", "-p"};
         TrainingResult result = TrainResnetWithCifar10.runExample(args);
-        Assert.assertNotNull(result);
 
-        Assert.assertTrue(result.getTrainEvaluation("Accuracy") >= 0.8f);
-        Assert.assertTrue(result.getValidateEvaluation("Accuracy") >= 0.68f);
-        Assert.assertTrue(result.getValidateLoss() < 1.1);
+        Assert.assertNotNull(result);
     }
 
     @Test
@@ -67,13 +48,13 @@ public class TrainResNetTest {
             throws ModelException, IOException, TranslateException {
         TestRequirements.linux();
         TestRequirements.nightly();
-        TestRequirements.gpu("MXNet");
+        TestRequirements.gpu("PyTorch", 4);
 
         // Limit max 4 gpu for cifar10 training to make it converge faster.
         // and only train 10 batch for unit test.
-        String[] args = {"-e", "10", "-g", "4", "--engine", "MXNet"};
+        String[] args = {"-e", "10", "-g", "4"};
 
-        Engine.getEngine("MXNet").setRandomSeed(SEED);
+        Engine.getEngine("PyTorch").setRandomSeed(SEED);
         TrainingResult result = TrainResnetWithCifar10.runExample(args);
         Assert.assertNotNull(result);
 
diff --git a/examples/src/test/java/ai/djl/examples/training/TrainSentimentAnalysisTest.java b/examples/src/test/java/ai/djl/examples/training/TrainSentimentAnalysisTest.java
index d726afdd4..c5ef28a55 100644
--- a/examples/src/test/java/ai/djl/examples/training/TrainSentimentAnalysisTest.java
+++ b/examples/src/test/java/ai/djl/examples/training/TrainSentimentAnalysisTest.java
@@ -27,7 +27,7 @@ public class TrainSentimentAnalysisTest {
             throws ModelException, TranslateException, IOException {
         TestRequirements.linux();
         TestRequirements.nightly();
-        TestRequirements.gpu("MXNet");
+        TestRequirements.gpu("MXNet", 1);
 
         String[] args = {"-e", "1", "-g", "1", "--engine", "MXNet"};
         TrainSentimentAnalysis.runExample(args);
diff --git a/examples/src/test/java/ai/djl/testing/TestRequirements.java b/examples/src/test/java/ai/djl/testing/TestRequirements.java
index 0e7d528cf..91aa7d4c8 100644
--- a/examples/src/test/java/ai/djl/testing/TestRequirements.java
+++ b/examples/src/test/java/ai/djl/testing/TestRequirements.java
@@ -14,6 +14,7 @@ package ai.djl.testing;
 
 import ai.djl.engine.Engine;
 import ai.djl.engine.EngineException;
+import ai.djl.util.cuda.CudaUtils;
 
 import org.testng.SkipException;
 
@@ -76,9 +77,16 @@ public final class TestRequirements {
     }
 
     /** Requires a test have at least one gpu. */
-    public static void gpu(String engine) {
-        if (Engine.getEngine(engine).getGpuCount() == 0) {
-            throw new SkipException("This test requires a GPU to run");
+    public static void gpu(String engine, int numGpu) {
+        if (Engine.getEngine(engine).getGpuCount() < numGpu) {
+            throw new SkipException("This test requires " + numGpu + " GPUs to run");
+        }
+    }
+
+    /** Avoid OOM on GPUs with multiple engines. */
+    public static void notGpu() {
+        if (CudaUtils.getGpuCount() > 0) {
+            throw new SkipException("This test requires CPU only machine to run");
         }
     }
 }
diff --git a/extensions/tokenizers/src/test/java/ai/djl/huggingface/tokenizers/TextEmbeddingTranslatorTest.java b/extensions/tokenizers/src/test/java/ai/djl/huggingface/tokenizers/TextEmbeddingTranslatorTest.java
index 4b8f883ff..3dddf301f 100644
--- a/extensions/tokenizers/src/test/java/ai/djl/huggingface/tokenizers/TextEmbeddingTranslatorTest.java
+++ b/extensions/tokenizers/src/test/java/ai/djl/huggingface/tokenizers/TextEmbeddingTranslatorTest.java
@@ -12,6 +12,7 @@
  */
 package ai.djl.huggingface.tokenizers;
 
+import ai.djl.Device;
 import ai.djl.Model;
 import ai.djl.ModelException;
 import ai.djl.huggingface.translator.TextEmbeddingTranslatorFactory;
@@ -104,6 +105,7 @@ public class TextEmbeddingTranslatorTest {
                         .optArgument("block_shapes", "(1,7,384)")
                         .optArgument("block_names", "last_hidden_state")
                         .optEngine("PyTorch")
+                        .optDevice(Device.cpu())
                         .optArgument("tokenizer", "bert-base-uncased")
                         .optArgument("pooling", "max")
                         .optOption("hasParameter", "false")
diff --git a/integration/src/main/java/ai/djl/integration/tests/training/listener/EarlyStoppingListenerTest.java b/integration/src/main/java/ai/djl/integration/tests/training/listener/EarlyStoppingListenerTest.java
index 9aee26614..15a6a2ae2 100644
--- a/integration/src/main/java/ai/djl/integration/tests/training/listener/EarlyStoppingListenerTest.java
+++ b/integration/src/main/java/ai/djl/integration/tests/training/listener/EarlyStoppingListenerTest.java
@@ -149,7 +149,7 @@ public class EarlyStoppingListenerTest {
                 }
 
                 TrainingResult trainingResult = trainer.getTrainingResult();
-                Assert.assertEquals(trainingResult.getEpoch(), 3);
+                Assert.assertTrue(trainingResult.getEpoch() < 5);
             }
         }
     }