Repository: deepjavalibrary/djl
Pull request #2070 (2022-10-08 20:48:01)

Transfer learning with pytorch engine on fresh fruit dataset
## Description ##
In this PR, the following features are enabled:
1. The transfer learning feature on PyTorch engine. User can opt in this feature by setting "retrain" to String "true" in `Criteria.builder().optOption(String key, String value)`, when loading the traced pretrained model.
2. The trainable variables(parameters) can have per variable loss function.



codecov-commenter: # [Codecov](https://codecov.io/gh/deepjavalibrary/djl/pull/2070?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) Report
Base: **72.08**% // Head: **70.18**% // Decreases project coverage by **`-1.90%`** :warning:
> Coverage data is based on head [(`74db9b9`)](https://codecov.io/gh/deepjavalibrary/djl/pull/2070?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) compared to base [(`bb5073f`)](https://codecov.io/gh/deepjavalibrary/djl/commit/bb5073f7ffc33f1883621fc1783dbb01e70b418f?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary).
> Patch coverage: 71.13% of modified lines in pull request are covered.

> :exclamation: Current head 74db9b9 differs from pull request most recent head ed0e7a6. Consider uploading reports for the commit ed0e7a6 to get more accurate results

<details><summary>Additional details and impacted files</summary>


```diff
@@             Coverage Diff              @@
##             master    #2070      +/-   ##
============================================
- Coverage     72.08%   70.18%   -1.91%     
- Complexity     5126     6084     +958     
============================================
  Files           473      606     +133     
  Lines         21970    26964    +4994     
  Branches       2351     2919     +568     
============================================
+ Hits          15838    18925    +3087     
- Misses         4925     6614    +1689     
- Partials       1207     1425     +218     
```


| [Impacted Files](https://codecov.io/gh/deepjavalibrary/djl/pull/2070?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) | Coverage Δ | |
|---|---|---|
| [api/src/main/java/ai/djl/modality/cv/Image.java](https://codecov.io/gh/deepjavalibrary/djl/pull/2070/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L0ltYWdlLmphdmE=) | `69.23% <ø> (-4.11%)` | :arrow_down: |
| [...rc/main/java/ai/djl/modality/cv/MultiBoxPrior.java](https://codecov.io/gh/deepjavalibrary/djl/pull/2070/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L011bHRpQm94UHJpb3IuamF2YQ==) | `76.00% <ø> (ø)` | |
| [...rc/main/java/ai/djl/modality/cv/output/Joints.java](https://codecov.io/gh/deepjavalibrary/djl/pull/2070/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L291dHB1dC9Kb2ludHMuamF2YQ==) | `71.42% <ø> (ø)` | |
| [.../main/java/ai/djl/modality/cv/output/Landmark.java](https://codecov.io/gh/deepjavalibrary/djl/pull/2070/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L291dHB1dC9MYW5kbWFyay5qYXZh) | `100.00% <ø> (ø)` | |
| [...main/java/ai/djl/modality/cv/output/Rectangle.java](https://codecov.io/gh/deepjavalibrary/djl/pull/2070/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L291dHB1dC9SZWN0YW5nbGUuamF2YQ==) | `72.41% <0.00%> (ø)` | |
| [...i/djl/modality/cv/translator/BigGANTranslator.java](https://codecov.io/gh/deepjavalibrary/djl/pull/2070/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3RyYW5zbGF0b3IvQmlnR0FOVHJhbnNsYXRvci5qYXZh) | `21.42% <0.00%> (-5.24%)` | :arrow_down: |
| [.../ai/djl/modality/cv/translator/YoloTranslator.java](https://codecov.io/gh/deepjavalibrary/djl/pull/2070/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3RyYW5zbGF0b3IvWW9sb1RyYW5zbGF0b3IuamF2YQ==) | `27.77% <0.00%> (+18.95%)` | :arrow_up: |
| [...modality/cv/translator/wrapper/FileTranslator.java](https://codecov.io/gh/deepjavalibrary/djl/pull/2070/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3RyYW5zbGF0b3Ivd3JhcHBlci9GaWxlVHJhbnNsYXRvci5qYXZh) | `44.44% <ø> (ø)` | |
| [...y/cv/translator/wrapper/InputStreamTranslator.java](https://codecov.io/gh/deepjavalibrary/djl/pull/2070/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3RyYW5zbGF0b3Ivd3JhcHBlci9JbnB1dFN0cmVhbVRyYW5zbGF0b3IuamF2YQ==) | `44.44% <ø> (ø)` | |
| [.../modality/cv/translator/wrapper/UrlTranslator.java](https://codecov.io/gh/deepjavalibrary/djl/pull/2070/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3RyYW5zbGF0b3Ivd3JhcHBlci9VcmxUcmFuc2xhdG9yLmphdmE=) | `44.44% <ø> (ø)` | |
| ... and [533 more](https://codecov.io/gh/deepjavalibrary/djl/pull/2070/diff?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) | |

Help us with your feedback. Take ten seconds to tell us [how you rate us](https://about.codecov.io/nps?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary). Have a feature suggestion? [Share it here.](https://app.codecov.io/gh/feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary)

</details>

[:umbrella: View full report at Codecov](https://codecov.io/gh/deepjavalibrary/djl/pull/2070?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary).   
:loudspeaker: Do you have feedback about the report comment? [Let us know in this issue](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary).


Patch:
diff --git a/api/src/main/java/ai/djl/modality/cv/transform/OneHot.java b/api/src/main/java/ai/djl/modality/cv/transform/OneHot.java
new file mode 100644
index 000000000..0e99c4603
--- /dev/null
+++ b/api/src/main/java/ai/djl/modality/cv/transform/OneHot.java
@@ -0,0 +1,37 @@
+/*
+ * Copyright 2022 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+ * with the License. A copy of the License is located at
+ *
+ * http://aws.amazon.com/apache2.0/
+ *
+ * or in the "license" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES
+ * OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
+ * and limitations under the License.
+ */
+package ai.djl.modality.cv.transform;
+
+import ai.djl.ndarray.NDArray;
+import ai.djl.translate.Transform;
+
+/** A {@link Transform} that converts the labels {@link NDArray} to one-hot labels. */
+public class OneHot implements Transform {
+
+    private int numClass;
+
+    /**
+     * Creates a {@code toOneHot} {@link Transform} that converts the sparse label to one-hot label.
+     *
+     * @param numClass number of classes
+     */
+    public OneHot(int numClass) {
+        this.numClass = numClass;
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public NDArray transform(NDArray array) {
+        return array.oneHot(numClass);
+    }
+}
diff --git a/api/src/main/java/ai/djl/modality/cv/transform/RandomFlipLeftRight.java b/api/src/main/java/ai/djl/modality/cv/transform/RandomFlipLeftRight.java
index 74b19c1e5..f83d220a9 100644
--- a/api/src/main/java/ai/djl/modality/cv/transform/RandomFlipLeftRight.java
+++ b/api/src/main/java/ai/djl/modality/cv/transform/RandomFlipLeftRight.java
@@ -12,18 +12,37 @@
  */
 package ai.djl.modality.cv.transform;
 
-import ai.djl.modality.cv.util.NDImageUtils;
 import ai.djl.ndarray.NDArray;
 import ai.djl.translate.Transform;
+import ai.djl.util.RandomUtils;
+
+import java.util.Random;
 
 /**
  * A {@link Transform} that randomly flip the input image left to right with a probability of 0.5.
  */
 public class RandomFlipLeftRight implements Transform {
+    Integer seed;
+
+    /** Creates a new instance of {@code RandomFlipLeftRight}. */
+    public RandomFlipLeftRight() {}
+
+    /**
+     * Creates a new instance of {@code RandomFlipLeftRight} with the given seed.
+     *
+     * @param seed the value of the seed
+     */
+    public RandomFlipLeftRight(int seed) {
+        this.seed = seed;
+    }
 
     /** {@inheritDoc} */
     @Override
     public NDArray transform(NDArray array) {
-        return NDImageUtils.randomFlipLeftRight(array);
+        Random rnd = (seed != null) ? new Random(seed) : RandomUtils.RANDOM;
+        if (rnd.nextFloat() > 0.5) {
+            array.flip(1);
+        }
+        return array;
     }
 }
diff --git a/api/src/main/java/ai/djl/modality/cv/transform/RandomFlipTopBottom.java b/api/src/main/java/ai/djl/modality/cv/transform/RandomFlipTopBottom.java
index 73c229b55..57a4b4fa0 100644
--- a/api/src/main/java/ai/djl/modality/cv/transform/RandomFlipTopBottom.java
+++ b/api/src/main/java/ai/djl/modality/cv/transform/RandomFlipTopBottom.java
@@ -12,18 +12,38 @@
  */
 package ai.djl.modality.cv.transform;
 
-import ai.djl.modality.cv.util.NDImageUtils;
 import ai.djl.ndarray.NDArray;
 import ai.djl.translate.Transform;
+import ai.djl.util.RandomUtils;
+
+import java.util.Random;
 
 /**
  * A {@link Transform} that randomly flip the input image top to bottom with a probability of 0.5.
  */
 public class RandomFlipTopBottom implements Transform {
 
+    Integer seed;
+
+    /** Creates a new instance of {@code RandomFlipTopBottom}. */
+    public RandomFlipTopBottom() {}
+
+    /**
+     * Creates a new instance of {@code RandomFlipTopBottom} with the given seed.
+     *
+     * @param seed the value of the seed
+     */
+    public RandomFlipTopBottom(int seed) {
+        this.seed = seed;
+    }
+
     /** {@inheritDoc} */
     @Override
     public NDArray transform(NDArray array) {
-        return NDImageUtils.randomFlipTopBottom(array);
+        Random rnd = (seed != null) ? new Random(seed) : RandomUtils.RANDOM;
+        if (rnd.nextFloat() > 0.5) {
+            array.flip(0);
+        }
+        return array;
     }
 }
diff --git a/api/src/main/java/ai/djl/modality/cv/transform/RandomResizedCrop.java b/api/src/main/java/ai/djl/modality/cv/transform/RandomResizedCrop.java
index b04dbc0b7..3f5cc1f03 100644
--- a/api/src/main/java/ai/djl/modality/cv/transform/RandomResizedCrop.java
+++ b/api/src/main/java/ai/djl/modality/cv/transform/RandomResizedCrop.java
@@ -51,6 +51,21 @@ public class RandomResizedCrop implements Transform {
         this.maxAspectRatio = maxAspectRatio;
     }
 
+    /**
+     * Creates a {@code RandomResizedCrop} {@link Transform}.
+     *
+     * @param width the output width of the image
+     * @param height the output height of the image
+     */
+    public RandomResizedCrop(int width, int height) {
+        this.width = width;
+        this.height = height;
+        this.minAreaScale = 0.08;
+        this.maxAreaScale = 1.0;
+        this.minAspectRatio = 3.0 / 4.0;
+        this.maxAspectRatio = 4.0 / 3.0;
+    }
+
     /** {@inheritDoc} */
     @Override
     public NDArray transform(NDArray array) {
diff --git a/api/src/main/java/ai/djl/modality/cv/util/NDImageUtils.java b/api/src/main/java/ai/djl/modality/cv/util/NDImageUtils.java
index 62d044b81..6d9714294 100644
--- a/api/src/main/java/ai/djl/modality/cv/util/NDImageUtils.java
+++ b/api/src/main/java/ai/djl/modality/cv/util/NDImageUtils.java
@@ -30,7 +30,7 @@ public final class NDImageUtils {
      *
      * @param image the image to resize
      * @param size the desired size
-     * @return the resized NDList
+     * @return the resized NDArray
      */
     public static NDArray resize(NDArray image, int size) {
         return resize(image, size, size, Image.Interpolation.BILINEAR);
@@ -42,7 +42,7 @@ public final class NDImageUtils {
      * @param image the image to resize
      * @param width the desired width
      * @param height the desired height
-     * @return the resized NDList
+     * @return the resized NDArray
      */
     public static NDArray resize(NDArray image, int width, int height) {
         return resize(image, width, height, Image.Interpolation.BILINEAR);
@@ -55,7 +55,7 @@ public final class NDImageUtils {
      * @param width the desired width
      * @param height the desired height
      * @param interpolation the desired interpolation
-     * @return the resized NDList
+     * @return the resized NDArray
      */
     public static NDArray resize(
             NDArray image, int width, int height, Image.Interpolation interpolation) {
@@ -334,6 +334,6 @@ public final class NDImageUtils {
         } else if (shape.get(2) == 1 || shape.get(2) == 3) {
             return false;
         }
-        throw new IllegalArgumentException("Image is not CHW or HWC");
+        throw new IllegalArgumentException("Image is neither CHW nor HWC");
     }
 }
diff --git a/api/src/main/java/ai/djl/nn/Block.java b/api/src/main/java/ai/djl/nn/Block.java
index 287bb768b..9bf52dcca 100644
--- a/api/src/main/java/ai/djl/nn/Block.java
+++ b/api/src/main/java/ai/djl/nn/Block.java
@@ -194,7 +194,8 @@ public interface Block {
     void initialize(NDManager manager, DataType dataType, Shape... inputShapes);
 
     /**
-     * Returns a boolean whether the block is initialized.
+     * Returns a boolean whether the block is initialized (block has inputShape and params have
+     * nonNull array).
      *
      * @return whether the block is initialized
      */
diff --git a/api/src/main/java/ai/djl/nn/Parameter.java b/api/src/main/java/ai/djl/nn/Parameter.java
index 857e4c474..2c36d3f11 100644
--- a/api/src/main/java/ai/djl/nn/Parameter.java
+++ b/api/src/main/java/ai/djl/nn/Parameter.java
@@ -113,6 +113,15 @@ public class Parameter implements AutoCloseable {
         this.shape = shape;
     }
 
+    /**
+     * Gets the shape of this {@code Parameter}.
+     *
+     * @return the shape of this {@code Parameter}
+     */
+    public Shape getShape() {
+        return shape;
+    }
+
     /**
      * Gets the values of this {@code Parameter} as an {@link NDArray}.
      *
@@ -170,6 +179,16 @@ public class Parameter implements AutoCloseable {
         this.initializer = initializer;
     }
 
+    /**
+     * Returns the {@link Initializer} for this {@code Parameter}, if not already set. If overwrite
+     * flag is true, sets the initializer regardless.
+     *
+     * @return the initializer of this {@code Parameter}
+     */
+    public Initializer getInitializer() {
+        return initializer;
+    }
+
     /**
      * Initializes the parameter with the given {@link NDManager}, with given {@link DataType} for
      * the given expected input shapes.
@@ -178,9 +197,13 @@ public class Parameter implements AutoCloseable {
      * @param dataType the datatype of the {@code Parameter}
      */
     public void initialize(NDManager manager, DataType dataType) {
-        Objects.requireNonNull(initializer, "No initializer has been set");
-        Objects.requireNonNull(shape, "No parameter shape has been set");
         if (!isInitialized()) {
+            // Params in a PtSymbolBlock is set during model loading and its isInitialized()=true.
+            // Shouldn't further initialize it.
+            Objects.requireNonNull(initializer, "No initializer has been set");
+            // Params in a PtSymbolBlock can have null shape, but are still initialized (has nonNull
+            // array)
+            Objects.requireNonNull(shape, "No parameter shape has been set");
             array = initializer.initialize(manager, shape, dataType);
             array.setName(name);
         }
diff --git a/api/src/main/java/ai/djl/nn/core/Linear.java b/api/src/main/java/ai/djl/nn/core/Linear.java
index 22003b525..21e7941ae 100644
--- a/api/src/main/java/ai/djl/nn/core/Linear.java
+++ b/api/src/main/java/ai/djl/nn/core/Linear.java
@@ -35,10 +35,19 @@ import java.util.Collections;
  * <p>It has the following shapes:
  *
  * <ul>
- *   <li>input X: [x1, x2, …, xn, input_dim]
+ *   <li>input X: [x1, x2, ..., xn, input_dim]
  *   <li>weight W: [units, input_dim]
  *   <li>Bias b: [units]
- *   <li>output Y: [x1, x2, …, xn, units]
+ *   <li>output Y: [x1, x2, ..., xn, units]
+ * </ul>
+ *
+ * <p>It is most typically used with a simple batched 1D input. In that case, the shape would be:
+ *
+ * <ul>
+ *   <li>input X: [batch_num, input_dim]
+ *   <li>weight W: [units, input_dim]
+ *   <li>Bias b: [units]
+ *   <li>output Y: [batch_num, units]
  * </ul>
  *
  * <p>The Linear block should be constructed using {@link Linear.Builder}.
diff --git a/api/src/main/java/ai/djl/training/evaluator/Accuracy.java b/api/src/main/java/ai/djl/training/evaluator/Accuracy.java
index 5b9695243..6bb6803ed 100644
--- a/api/src/main/java/ai/djl/training/evaluator/Accuracy.java
+++ b/api/src/main/java/ai/djl/training/evaluator/Accuracy.java
@@ -61,7 +61,8 @@ public class Accuracy extends AbstractAccuracy {
             predictionReduced = predictionReduced.reshape(label.getShape());
         } else {
             // Multi-class, one-hot label
-            predictionReduced = prediction;
+            predictionReduced = prediction.argMax(axis);
+            label = label.argMax(axis);
         }
         // result of sum is int64 now
         long total = label.size();
diff --git a/api/src/main/java/ai/djl/training/loss/MaskedSoftmaxCrossEntropyLoss.java b/api/src/main/java/ai/djl/training/loss/MaskedSoftmaxCrossEntropyLoss.java
index 1d90da598..3444509f6 100644
--- a/api/src/main/java/ai/djl/training/loss/MaskedSoftmaxCrossEntropyLoss.java
+++ b/api/src/main/java/ai/djl/training/loss/MaskedSoftmaxCrossEntropyLoss.java
@@ -47,9 +47,10 @@ public class MaskedSoftmaxCrossEntropyLoss extends Loss {
      * @param name the name of the loss
      * @param weight the weight to apply on the loss value, default 1
      * @param classAxis the axis that represents the class probabilities, default -1
-     * @param sparseLabel whether labels are integer array or probabilities, default true
-     * @param fromLogit whether predictions are log probabilities or un-normalized numbers, default
-     *     false
+     * @param sparseLabel whether labels are 1-D integer array of [batch_size] (false) or 2-D
+     *     probabilities of [batch_size, n-class] (true), default true
+     * @param fromLogit if true, the inputs are assumed to be the numbers before being applied with
+     *     softmax. Then logSoftmax will be applied to input, default false
      */
     public MaskedSoftmaxCrossEntropyLoss(
             String name, float weight, int classAxis, boolean sparseLabel, boolean fromLogit) {
diff --git a/api/src/main/java/ai/djl/training/loss/SoftmaxCrossEntropyLoss.java b/api/src/main/java/ai/djl/training/loss/SoftmaxCrossEntropyLoss.java
index 0bc4a6fc5..0269cb7ea 100644
--- a/api/src/main/java/ai/djl/training/loss/SoftmaxCrossEntropyLoss.java
+++ b/api/src/main/java/ai/djl/training/loss/SoftmaxCrossEntropyLoss.java
@@ -43,6 +43,8 @@ public class SoftmaxCrossEntropyLoss extends Loss {
      * @param name the name of the loss
      */
     public SoftmaxCrossEntropyLoss(String name) {
+        // By default, fromLogit=true, means it takes the prediction before being
+        // applied softmax.
         this(name, 1, -1, true, true);
     }
 
@@ -52,10 +54,10 @@ public class SoftmaxCrossEntropyLoss extends Loss {
      * @param name the name of the loss
      * @param weight the weight to apply on the loss value, default 1
      * @param classAxis the axis that represents the class probabilities, default -1
-     * @param sparseLabel whether labels are 1-D integer array or 2-D probabilities of [batch_size,
-     *     n-class], default true
-     * @param fromLogit whether predictions are un-normalized numbers or log probabilities, if true,
-     *     logSoftmax will be applied to input, default true
+     * @param sparseLabel whether labels are rank-1 integer array of [batch_size] (false) or rank-2
+     *     one-hot of [batch_size, n-class] (true), default true
+     * @param fromLogit if true, the inputs are assumed to be the numbers before being applied with
+     *     softmax. Then logSoftmax will be applied to input, default true
      */
     public SoftmaxCrossEntropyLoss(
             String name, float weight, int classAxis, boolean sparseLabel, boolean fromLogit) {
diff --git a/api/src/main/java/ai/djl/training/optimizer/Adagrad.java b/api/src/main/java/ai/djl/training/optimizer/Adagrad.java
index 7b125a385..5a5ece17f 100644
--- a/api/src/main/java/ai/djl/training/optimizer/Adagrad.java
+++ b/api/src/main/java/ai/djl/training/optimizer/Adagrad.java
@@ -17,6 +17,7 @@ import ai.djl.ndarray.NDArray;
 import ai.djl.ndarray.NDList;
 import ai.djl.ndarray.internal.NDArrayEx;
 import ai.djl.ndarray.types.SparseFormat;
+import ai.djl.training.tracker.ParameterTracker;
 import ai.djl.training.tracker.Tracker;
 
 import java.util.Map;
@@ -40,7 +41,7 @@ import java.util.concurrent.ConcurrentHashMap;
  */
 public class Adagrad extends Optimizer {
 
-    private Tracker learningRateTracker;
+    private ParameterTracker learningRateTracker;
     private float epsilon;
 
     private Map<String, Map<Device, NDArray>> history;
@@ -61,7 +62,7 @@ public class Adagrad extends Optimizer {
     @Override
     public void update(String parameterId, NDArray weight, NDArray grad) {
         int t = updateCount(parameterId);
-        float newLearningRate = learningRateTracker.getNewValue(t);
+        float newLearningRate = learningRateTracker.getNewValue(parameterId, t);
         float weightDecay = getWeightDecay();
 
         if (Float.isNaN(newLearningRate)
@@ -98,7 +99,7 @@ public class Adagrad extends Optimizer {
     /** The Builder to construct an {@link Adagrad} object. */
     public static final class Builder extends OptimizerBuilder<Builder> {
 
-        private Tracker learningRateTracker = Tracker.fixed(0.001f);
+        private ParameterTracker learningRateTracker = Tracker.fixed(0.001f);
         private float epsilon = 1e-8f;
 
         Builder() {}
@@ -110,12 +111,12 @@ public class Adagrad extends Optimizer {
         }
 
         /**
-         * Sets the {@link Tracker} for this optimizer.
+         * Sets the {@link ParameterTracker} for this optimizer.
          *
-         * @param learningRateTracker the {@link Tracker} to be set
+         * @param learningRateTracker the {@link ParameterTracker} to be set
          * @return this {@code Builder}
          */
-        public Builder optLearningRateTracker(Tracker learningRateTracker) {
+        public Builder optLearningRateTracker(ParameterTracker learningRateTracker) {
             this.learningRateTracker = learningRateTracker;
             return this;
         }
diff --git a/api/src/main/java/ai/djl/training/optimizer/Adam.java b/api/src/main/java/ai/djl/training/optimizer/Adam.java
index e4544f84a..5232dbb50 100644
--- a/api/src/main/java/ai/djl/training/optimizer/Adam.java
+++ b/api/src/main/java/ai/djl/training/optimizer/Adam.java
@@ -16,6 +16,7 @@ import ai.djl.Device;
 import ai.djl.ndarray.NDArray;
 import ai.djl.ndarray.NDList;
 import ai.djl.ndarray.internal.NDArrayEx;
+import ai.djl.training.tracker.ParameterTracker;
 import ai.djl.training.tracker.Tracker;
 import ai.djl.util.Preconditions;
 
@@ -38,7 +39,7 @@ import java.util.concurrent.ConcurrentHashMap;
  */
 public class Adam extends Optimizer {
 
-    private Tracker learningRateTracker;
+    private ParameterTracker learningRateTracker;
     private float beta1;
     private float beta2;
     private float epsilon;
@@ -67,7 +68,7 @@ public class Adam extends Optimizer {
         int t = updateCount(parameterId);
         double coef1 = 1.0 - Math.pow(beta1, t);
         double coef2 = 1.0 - Math.pow(beta2, t);
-        float lr = learningRateTracker.getNewValue(t);
+        float lr = learningRateTracker.getNewValue(parameterId, t);
         float newLearningRate = (float) (lr * Math.sqrt(coef2) / coef1);
         float weightDecay = getWeightDecay();
 
@@ -117,7 +118,7 @@ public class Adam extends Optimizer {
     /** The Builder to construct an {@link Adam} object. */
     public static final class Builder extends OptimizerBuilder<Builder> {
 
-        private Tracker learningRateTracker = Tracker.fixed(0.001f);
+        private ParameterTracker learningRateTracker = Tracker.fixed(0.001f);
         private float beta1 = 0.9f;
         private float beta2 = 0.999f;
         private float epsilon = 1e-8f;
@@ -131,12 +132,12 @@ public class Adam extends Optimizer {
         }
 
         /**
-         * Sets the {@link Tracker} for this optimizer.
+         * Sets the {@link ParameterTracker} for this optimizer.
          *
-         * @param learningRateTracker the {@link Tracker} to be set
+         * @param learningRateTracker the {@link ParameterTracker} to be set
          * @return this {@code Builder}
          */
-        public Builder optLearningRateTracker(Tracker learningRateTracker) {
+        public Builder optLearningRateTracker(ParameterTracker learningRateTracker) {
             this.learningRateTracker = learningRateTracker;
             return this;
         }
diff --git a/api/src/main/java/ai/djl/training/optimizer/Nag.java b/api/src/main/java/ai/djl/training/optimizer/Nag.java
index b542b3fc1..be64f521a 100644
--- a/api/src/main/java/ai/djl/training/optimizer/Nag.java
+++ b/api/src/main/java/ai/djl/training/optimizer/Nag.java
@@ -17,7 +17,7 @@ import ai.djl.Device;
 import ai.djl.ndarray.NDArray;
 import ai.djl.ndarray.NDList;
 import ai.djl.ndarray.internal.NDArrayEx;
-import ai.djl.training.tracker.Tracker;
+import ai.djl.training.tracker.ParameterTracker;
 import ai.djl.util.Preconditions;
 
 import java.util.Map;
@@ -33,7 +33,7 @@ import java.util.concurrent.ConcurrentHashMap;
  */
 public class Nag extends Optimizer {
 
-    private Tracker learningRateTracker;
+    private ParameterTracker learningRateTracker;
     private float momentum;
     private Map<String, Map<Device, NDArray>> momentumStates;
 
@@ -53,7 +53,8 @@ public class Nag extends Optimizer {
     @Override
     public void update(String parameterId, NDArray weight, NDArray grad) {
         // TODO: Support Mixed precision Sparse
-        float newLearningRate = learningRateTracker.getNewValue(updateCount(parameterId));
+        float newLearningRate =
+                learningRateTracker.getNewValue(parameterId, updateCount(parameterId));
         float weightDecay = getWeightDecay();
         NDList inputs;
         if (momentum != 0f) {
@@ -79,18 +80,18 @@ public class Nag extends Optimizer {
     /** The Builder to construct an {@link Nag} object. */
     public static final class Builder extends OptimizerBuilder<Builder> {
 
-        Tracker learningRateTracker;
+        ParameterTracker learningRateTracker;
         float momentum;
 
         Builder() {}
 
         /**
-         * Sets the {@link Tracker} for this optimizer.
+         * Sets the {@link ParameterTracker} for this optimizer.
          *
-         * @param learningRateTracker the {@link Tracker} to be set
+         * @param learningRateTracker the {@link ParameterTracker} to be set
          * @return this {@code Builder}
          */
-        public Builder setLearningRateTracker(Tracker learningRateTracker) {
+        public Builder setLearningRateTracker(ParameterTracker learningRateTracker) {
             this.learningRateTracker = learningRateTracker;
             return this;
         }
diff --git a/api/src/main/java/ai/djl/training/optimizer/RmsProp.java b/api/src/main/java/ai/djl/training/optimizer/RmsProp.java
index 9aade4129..5c5909adf 100644
--- a/api/src/main/java/ai/djl/training/optimizer/RmsProp.java
+++ b/api/src/main/java/ai/djl/training/optimizer/RmsProp.java
@@ -16,6 +16,7 @@ import ai.djl.Device;
 import ai.djl.ndarray.NDArray;
 import ai.djl.ndarray.NDList;
 import ai.djl.ndarray.internal.NDArrayEx;
+import ai.djl.training.tracker.ParameterTracker;
 import ai.djl.training.tracker.Tracker;
 
 import java.util.Map;
@@ -55,7 +56,7 @@ import java.util.concurrent.ConcurrentHashMap;
  */
 public class RmsProp extends Optimizer {
 
-    private Tracker learningRateTracker;
+    private ParameterTracker learningRateTracker;
     private float rho;
     private float momentum;
     private float epsilon;
@@ -85,7 +86,8 @@ public class RmsProp extends Optimizer {
     /** {@inheritDoc} */
     @Override
     public void update(String parameterId, NDArray weight, NDArray grad) {
-        float newLearningRate = learningRateTracker.getNewValue(updateCount(parameterId));
+        float newLearningRate =
+                learningRateTracker.getNewValue(parameterId, updateCount(parameterId));
         float weightDecay = getWeightDecay();
 
         if (Float.isNaN(newLearningRate)
@@ -159,7 +161,7 @@ public class RmsProp extends Optimizer {
     /** The Builder to construct an {@link RmsProp} object. */
     public static final class Builder extends OptimizerBuilder<Builder> {
 
-        private Tracker learningRateTracker = Tracker.fixed(0.001f);
+        private ParameterTracker learningRateTracker = Tracker.fixed(0.001f);
         private float rho = 0.9f;
         private float momentum = 0.9f;
         private float epsilon = 1e-8f;
@@ -174,12 +176,12 @@ public class RmsProp extends Optimizer {
         }
 
         /**
-         * Sets the {@link Tracker} for this optimizer.
+         * Sets the {@link ParameterTracker} for this optimizer.
          *
-         * @param learningRateTracker the {@link Tracker} to be set
+         * @param learningRateTracker the {@link ParameterTracker} to be set
          * @return this {@code Builder}
          */
-        public Builder optLearningRateTracker(Tracker learningRateTracker) {
+        public Builder optLearningRateTracker(ParameterTracker learningRateTracker) {
             this.learningRateTracker = learningRateTracker;
             return this;
         }
diff --git a/api/src/main/java/ai/djl/training/optimizer/Sgd.java b/api/src/main/java/ai/djl/training/optimizer/Sgd.java
index d7ed0eacd..4e5f706bd 100644
--- a/api/src/main/java/ai/djl/training/optimizer/Sgd.java
+++ b/api/src/main/java/ai/djl/training/optimizer/Sgd.java
@@ -16,7 +16,7 @@ import ai.djl.Device;
 import ai.djl.ndarray.NDArray;
 import ai.djl.ndarray.NDList;
 import ai.djl.ndarray.internal.NDArrayEx;
-import ai.djl.training.tracker.Tracker;
+import ai.djl.training.tracker.ParameterTracker;
 
 import java.util.Map;
 import java.util.Objects;
@@ -37,7 +37,7 @@ import java.util.concurrent.ConcurrentHashMap;
  */
 public class Sgd extends Optimizer {
 
-    private Tracker learningRateTracker;
+    private ParameterTracker learningRateTracker;
     private float momentum;
     private Map<String, Map<Device, NDArray>> momentumStates;
 
@@ -58,7 +58,7 @@ public class Sgd extends Optimizer {
     public void update(String parameterId, NDArray weight, NDArray grad) {
         // TODO: Support Mixed precision Sparse
         float weightDecay = getWeightDecay();
-        float learningRate = learningRateTracker.getNewValue(updateCount(parameterId));
+        float learningRate = learningRateTracker.getNewValue(parameterId, updateCount(parameterId));
         NDList inputs;
         if (momentum != 0f) {
             NDArray state =
@@ -81,7 +81,7 @@ public class Sgd extends Optimizer {
     /** The Builder to construct an {@link Sgd} object. */
     public static final class Builder extends OptimizerBuilder<Builder> {
 
-        Tracker learningRateTracker;
+        ParameterTracker learningRateTracker;
         float momentum;
 
         Builder() {}
@@ -93,12 +93,12 @@ public class Sgd extends Optimizer {
         }
 
         /**
-         * Sets the {@link Tracker} for this optimizer.
+         * Sets the {@link ParameterTracker} for this optimizer.
          *
-         * @param learningRateTracker the {@link Tracker} to be set
+         * @param learningRateTracker the {@link ParameterTracker} to be set
          * @return this {@code Builder}
          */
-        public Builder setLearningRateTracker(Tracker learningRateTracker) {
+        public Builder setLearningRateTracker(ParameterTracker learningRateTracker) {
             this.learningRateTracker = learningRateTracker;
             return this;
         }
diff --git a/api/src/main/java/ai/djl/training/tracker/FixedPerVarTracker.java b/api/src/main/java/ai/djl/training/tracker/FixedPerVarTracker.java
new file mode 100644
index 000000000..87c85dd1a
--- /dev/null
+++ b/api/src/main/java/ai/djl/training/tracker/FixedPerVarTracker.java
@@ -0,0 +1,105 @@
+/*
+ * Copyright 2022 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+ * with the License. A copy of the License is located at
+ *
+ * http://aws.amazon.com/apache2.0/
+ *
+ * or in the "license" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES
+ * OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
+ * and limitations under the License.
+ */
+package ai.djl.training.tracker;
+
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+
+/**
+ * {@link FixedPerVarTracker} is an implementation of {@link Tracker} which returns a fixed value.
+ *
+ * @see Tracker
+ */
+public class FixedPerVarTracker implements ParameterTracker {
+
+    private float value;
+    private Map<String, Float> valueMap;
+
+    /**
+     * Creates a new instance of {@link FixedPerVarTracker}.
+     *
+     * @param builder the builder used to build this object
+     */
+    public FixedPerVarTracker(Builder builder) {
+        this.value = builder.value;
+        this.valueMap = builder.valueMap;
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public float getNewValue(String parameterId, int numUpdate) {
+        return valueMap.getOrDefault(parameterId, this.value);
+    }
+
+    /**
+     * Creates a builder to build a {@link FixedPerVarTracker}.
+     *
+     * @return a new builder
+     */
+    public static Builder builder() {
+        return new Builder();
+    }
+
+    /** The Builder to construct an {@link FixedPerVarTracker} object. */
+    public static final class Builder {
+
+        private float value;
+        private Map<String, Float> valueMap = new ConcurrentHashMap<>();
+
+        /** Create a builder for {@link FixedPerVarTracker}. */
+        private Builder() {}
+
+        /**
+         * Set the default learning rate.
+         *
+         * @param value the default learning rate
+         * @return builder
+         */
+        public Builder setDefaultValue(float value) {
+            this.value = value;
+            return this;
+        }
+
+        /**
+         * Add a kv pair of parameter and its learning rate.
+         *
+         * @param parameterId the parameter id
+         * @param value the default learning rate
+         * @return builder
+         */
+        public Builder put(String parameterId, float value) {
+            this.valueMap.put(parameterId, value);
+            return this;
+        }
+
+        /**
+         * Add kv pairs of parameter and its learning rate.
+         *
+         * @param valueMap stores parameterId and learning rate
+         * @return builder
+         */
+        public Builder putAll(Map<String, Float> valueMap) {
+            this.valueMap.putAll(valueMap);
+            return this;
+        }
+
+        /**
+         * Builds a {@link FixedPerVarTracker} block.
+         *
+         * @return the {@link FixedPerVarTracker} block
+         */
+        public FixedPerVarTracker build() {
+            return new FixedPerVarTracker(this);
+        }
+    }
+}
diff --git a/api/src/main/java/ai/djl/training/tracker/ParameterTracker.java b/api/src/main/java/ai/djl/training/tracker/ParameterTracker.java
new file mode 100644
index 000000000..c09e91b5e
--- /dev/null
+++ b/api/src/main/java/ai/djl/training/tracker/ParameterTracker.java
@@ -0,0 +1,31 @@
+/*
+ * Copyright 2022 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+ * with the License. A copy of the License is located at
+ *
+ * http://aws.amazon.com/apache2.0/
+ *
+ * or in the "license" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES
+ * OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
+ * and limitations under the License.
+ */
+package ai.djl.training.tracker;
+
+/**
+ * A {@code Tracker} represents a collection of hyperparameters or {@link Tracker}s that changes
+ * gradually through the training process.
+ *
+ * @see Tracker
+ */
+public interface ParameterTracker {
+
+    /**
+     * Fetches the value after the given number of steps/updates for the parameter.
+     *
+     * @param parameterId the id of the parameter to get the new value for
+     * @param numUpdate the total number of steps/updates
+     * @return this {@code Builder}
+     */
+    float getNewValue(String parameterId, int numUpdate);
+}
diff --git a/api/src/main/java/ai/djl/training/tracker/Tracker.java b/api/src/main/java/ai/djl/training/tracker/Tracker.java
index 913765330..1f6f28188 100644
--- a/api/src/main/java/ai/djl/training/tracker/Tracker.java
+++ b/api/src/main/java/ai/djl/training/tracker/Tracker.java
@@ -13,13 +13,13 @@
 package ai.djl.training.tracker;
 
 /**
- * A {@code Tracker} represents a hyper-parameter that changes gradually through the training
+ * A {@code Tracker} represents a hyperparameter that changes gradually through the training
  * process.
  *
  * @see <a href="https://d2l.djl.ai/chapter_optimization/lr-scheduler.html">For tracking learning
  *     rates, the D2L chapter on learning rate scheduling</a>
  */
-public interface Tracker {
+public interface Tracker extends ParameterTracker {
 
     /**
      * Fetches the value after the given number of steps/updates.
@@ -29,6 +29,12 @@ public interface Tracker {
      */
     float getNewValue(int numUpdate);
 
+    /** {@inheritDoc} */
+    @Override
+    default float getNewValue(String parameterId, int numUpdate) {
+        return getNewValue(numUpdate);
+    }
+
     /**
      * Returns a new instance of {@link ai.djl.training.tracker.FactorTracker.Builder} that can
      * build an {@link FactorTracker}.
@@ -84,7 +90,7 @@ public interface Tracker {
      * Returns a new instance of {@link Tracker} with a fixed value.
      *
      * @param value the fixed value
-     * @return a instance of {@link Tracker} with a fixed value
+     * @return an instance of {@link Tracker} with a fixed value
      */
     static Tracker fixed(float value) {
         return FixedTracker.builder().setValue(value).build();
diff --git a/api/src/main/java/ai/djl/translate/StackBatchifier.java b/api/src/main/java/ai/djl/translate/StackBatchifier.java
index 8143d184d..d4eb46950 100644
--- a/api/src/main/java/ai/djl/translate/StackBatchifier.java
+++ b/api/src/main/java/ai/djl/translate/StackBatchifier.java
@@ -77,7 +77,11 @@ public class StackBatchifier implements Batchifier {
                     NDArray currInput = input.get(i);
                     if (!currInput.getShape().equals(kindDataShape)) {
                         throw new IllegalArgumentException(
-                                "You cannot batch data with different input shapes", e);
+                                "You cannot batch data with different input shapes"
+                                        + currInput.getShape()
+                                        + " vs "
+                                        + kindDataShape,
+                                e);
                     }
                     if (!currInput.getDataType().equals(kindDataType)) {
                         throw new IllegalArgumentException(
diff --git a/basicdataset/src/main/java/ai/djl/basicdataset/cv/classification/FruitsFreshAndRotten.java b/basicdataset/src/main/java/ai/djl/basicdataset/cv/classification/FruitsFreshAndRotten.java
new file mode 100644
index 000000000..cdc8be7e4
--- /dev/null
+++ b/basicdataset/src/main/java/ai/djl/basicdataset/cv/classification/FruitsFreshAndRotten.java
@@ -0,0 +1,225 @@
+/*
+ * Copyright 2022 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+ * with the License. A copy of the License is located at
+ *
+ * http://aws.amazon.com/apache2.0/
+ *
+ * or in the "license" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES
+ * OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
+ * and limitations under the License.
+ */
+package ai.djl.basicdataset.cv.classification;
+
+import ai.djl.Application;
+import ai.djl.basicdataset.BasicDatasets;
+import ai.djl.modality.cv.transform.ToTensor;
+import ai.djl.repository.Artifact;
+import ai.djl.repository.MRL;
+import ai.djl.repository.Repository;
+import ai.djl.translate.Pipeline;
+import ai.djl.util.Progress;
+
+import java.io.File;
+import java.io.IOException;
+import java.nio.file.Path;
+import java.nio.file.Paths;
+import java.util.Arrays;
+
+/**
+ * FruitRottenFresh classification dataset that contains the same fruit where rotten and fresh class
+ * are stored in different sub folders.
+ *
+ * <pre>
+ * It is structured similar to ImageFolders as follows:
+ *      root/freshapples/1.png
+ *      root/freshapples/2.png
+ *      ...
+ *      root/rottenapples/1.png
+ *      root/rottenapples/2.png
+ *      ...
+ *      root/freshbanana/1.png
+ *      root/freshbanana/2.png
+ *      ...
+ *      root/rottenbanana/1.png
+ *      root/rottenbanana/2.png
+ *      ...
+ * </pre>
+ */
+public final class FruitsFreshAndRotten extends AbstractImageFolder {
+
+    private static final String VERSION = "1.0";
+    private static final String ARTIFACT_ID = "fruit";
+
+    private MRL mrl;
+    private boolean prepared;
+
+    private FruitsFreshAndRotten(Builder builder) {
+        super(builder);
+        mrl = builder.getMrl();
+    }
+
+    /**
+     * Creates a new builder to build a {@link FruitsFreshAndRotten}.
+     *
+     * @return a new builder
+     */
+    public static Builder builder() {
+        return new Builder();
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    protected Path getImagePath(String key) {
+        return Paths.get(key);
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public void prepare(Progress progress) throws IOException {
+        // Use the code in ImageFolder
+        if (!prepared) {
+            mrl.prepare(null, progress);
+            loadSynset();
+            Path root = Paths.get(mrl.getRepository().getBaseUri());
+            if (progress != null) {
+                progress.reset("Preparing", 2);
+                progress.start(0);
+                listImages(root, synset);
+                progress.end();
+            } else {
+                listImages(root, synset);
+            }
+            prepared = true;
+        }
+    }
+
+    private void loadSynset() {
+        File root = new File(mrl.getRepository().getBaseUri());
+        File[] dir = root.listFiles(f -> f.isDirectory() && !f.getName().startsWith("."));
+        if (dir == null || dir.length == 0) {
+            throw new IllegalArgumentException(root + " not found or didn't have any file in it");
+        }
+        Arrays.sort(dir);
+        for (File file : dir) {
+            synset.add(file.getName());
+        }
+    }
+
+    /** A builder for the {@link FruitsFreshAndRotten}. */
+    public static final class Builder extends ImageFolderBuilder<Builder> {
+
+        String groupId;
+        String artifactId;
+        Usage usage;
+        private Repository optRepository;
+
+        /** Constructs a new builder. */
+        Builder() {
+            repository = BasicDatasets.REPOSITORY;
+            groupId = BasicDatasets.GROUP_ID;
+            artifactId = ARTIFACT_ID;
+            usage = Usage.TRAIN;
+        }
+
+        /** {@inheritDoc} */
+        @Override
+        public Builder self() {
+            return this;
+        }
+
+        /**
+         * Sets the optional usage.
+         *
+         * @param usage the usage
+         * @return this builder
+         */
+        public Builder optUsage(Usage usage) {
+            this.usage = usage;
+            return self();
+        }
+
+        /**
+         * Sets the optional repository.
+         *
+         * @param repository the repository
+         * @return this builder
+         */
+        public Builder optRepository(Repository repository) {
+            this.optRepository = repository;
+            return self();
+        }
+
+        /**
+         * Sets optional groupId.
+         *
+         * @param groupId the groupId}
+         * @return this builder
+         */
+        public Builder optGroupId(String groupId) {
+            this.groupId = groupId;
+            return this;
+        }
+
+        /**
+         * Sets the optional artifactId.
+         *
+         * @param artifactId the artifactId
+         * @return this builder
+         */
+        public Builder optArtifactId(String artifactId) {
+            if (artifactId.contains(":")) {
+                String[] tokens = artifactId.split(":");
+                groupId = tokens[0];
+                this.artifactId = tokens[1];
+            } else {
+                this.artifactId = artifactId;
+            }
+            return this;
+        }
+
+        /**
+         * Builds the {@link FruitsFreshAndRotten}.
+         *
+         * @return the {@link FruitsFreshAndRotten}
+         * @throws IOException if there is an issue
+         */
+        public FruitsFreshAndRotten build() throws IOException {
+            if (pipeline == null) {
+                pipeline = new Pipeline(new ToTensor());
+            }
+
+            if (optRepository != null) {
+                repository = optRepository;
+            } else {
+                MRL mrl = getMrl();
+                Artifact artifact = mrl.getDefaultArtifact();
+                // Downloading the cache happens here
+                mrl.prepare(artifact, null);
+
+                Artifact.Item item;
+                switch (usage) {
+                    case TRAIN:
+                        item = artifact.getFiles().get("train");
+                        break;
+                    case TEST:
+                        item = artifact.getFiles().get("test");
+                        break;
+                    case VALIDATION:
+                    default:
+                        throw new IOException("Only training and testing dataset supported.");
+                }
+                Path root = mrl.getRepository().getFile(item, "").toAbsolutePath();
+
+                // set repository
+                repository = Repository.newInstance("banana", root);
+            }
+            return new FruitsFreshAndRotten(this);
+        }
+
+        MRL getMrl() {
+            return repository.dataset(Application.CV.ANY, groupId, artifactId, VERSION);
+        }
+    }
+}
diff --git a/basicdataset/src/main/java/ai/djl/basicdataset/cv/classification/ImageFolder.java b/basicdataset/src/main/java/ai/djl/basicdataset/cv/classification/ImageFolder.java
index df5e887a4..8f242f39d 100644
--- a/basicdataset/src/main/java/ai/djl/basicdataset/cv/classification/ImageFolder.java
+++ b/basicdataset/src/main/java/ai/djl/basicdataset/cv/classification/ImageFolder.java
@@ -51,7 +51,7 @@ import java.util.Arrays;
  * // set the image folder path
  * Repository repository = Repository.newInstance("folder", Paths.get("/path/to/imagefolder/root");
  * ImageFolder dataset =
- *     new ImageFolder.Builder()
+ *     ImageFolder.builder()
  *         .setRepository(repository)
  *         .addTransform(new Resize(100, 100)) // Use image transforms as necessary for your data
  *         .addTransform(new ToTensor()) // Usually required as the last transform to convert images to tensors
diff --git a/engines/mxnet/mxnet-engine/src/main/java/ai/djl/mxnet/engine/MxModel.java b/engines/mxnet/mxnet-engine/src/main/java/ai/djl/mxnet/engine/MxModel.java
index fe42c2f7a..88289b2aa 100644
--- a/engines/mxnet/mxnet-engine/src/main/java/ai/djl/mxnet/engine/MxModel.java
+++ b/engines/mxnet/mxnet-engine/src/main/java/ai/djl/mxnet/engine/MxModel.java
@@ -107,7 +107,11 @@ public class MxModel extends BaseModel {
             paramFile = paramPathResolver(prefix, options);
             if (paramFile == null && block == null) {
                 throw new FileNotFoundException(
-                        "Parameter file with prefix: " + prefix + " not found in: " + modelDir);
+                        "Parameter file with prefix: "
+                                + prefix
+                                + " not found in: "
+                                + modelDir
+                                + " or not readable by the engine.");
             }
         }
 
diff --git a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtModel.java b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtModel.java
index 9c28885a6..49f1a1f6a 100644
--- a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtModel.java
+++ b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtModel.java
@@ -78,12 +78,14 @@ public class PtModel extends BaseModel {
             String[] extraFileKeys = new String[0];
             String[] extraFileValues = new String[0];
             boolean mapLocation = false;
+            boolean retrain = false;
             // load jit extra files
             if (options != null) {
                 if (options.containsKey("extraFiles")) {
                     extraFileKeys = ((String) options.get("extraFiles")).split(",");
                     extraFileValues = new String[extraFileKeys.length];
                 }
+                retrain = Boolean.parseBoolean((String) options.get("retrain"));
                 mapLocation = Boolean.parseBoolean((String) options.get("mapLocation"));
             }
             block =
@@ -92,10 +94,16 @@ public class PtModel extends BaseModel {
                             modelFile,
                             mapLocation,
                             extraFileKeys,
-                            extraFileValues);
+                            extraFileValues,
+                            retrain);
             for (int i = 0; i < extraFileKeys.length; i++) {
                 properties.put(extraFileKeys[i], extraFileValues[i]);
             }
+            // By default, the parameters are frozen, since before adding this training feature
+            // `retrain`, it was frozen already by setting `JITCallGuard guard`, which disables
+            // autograd. Also, the pretrained parameters usually should not be updated too much. It
+            // is safe to freeze it. Users may unfreeze it and set their learning rate small.
+            block.freezeParameters(!retrain);
         } else {
             boolean hasParameter = true;
             if (options != null) {
diff --git a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtNDArrayEx.java b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtNDArrayEx.java
index c22a5347a..28375ef38 100644
--- a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtNDArrayEx.java
+++ b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtNDArrayEx.java
@@ -575,6 +575,7 @@ public class PtNDArrayEx implements NDArrayEx {
             if (dim == 3) {
                 result = result.expandDims(0);
             }
+            // change from HWC to CHW order
             result = result.transpose(0, 3, 1, 2);
             result =
                     JniUtils.interpolate(
diff --git a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/JniUtils.java b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/JniUtils.java
index 192b1c444..a4bb78cb6 100644
--- a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/JniUtils.java
+++ b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/JniUtils.java
@@ -1542,7 +1542,8 @@ public final class JniUtils {
             Path path,
             boolean mapLocation,
             String[] extraFileKeys,
-            String[] extraFileValues) {
+            String[] extraFileValues,
+            boolean retrain) {
         Device device = manager.getDevice();
         // MPS doesn't support mapLocation
         if ("mps".equals(device.getDeviceType())) {
@@ -1556,7 +1557,8 @@ public final class JniUtils {
                         new int[] {PtDeviceType.toDeviceType(device), device.getDeviceId()},
                         mapLocation,
                         extraFileKeys,
-                        extraFileValues);
+                        extraFileValues,
+                        retrain);
         return new PtSymbolBlock(manager, handle);
     }
 
diff --git a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/PyTorchLibrary.java b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/PyTorchLibrary.java
index 57ba028e4..378efeae4 100644
--- a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/PyTorchLibrary.java
+++ b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/PyTorchLibrary.java
@@ -491,7 +491,8 @@ final class PyTorchLibrary {
             int[] device,
             boolean mapLocation,
             String[] extraFileNames,
-            String[] extraFileValues);
+            String[] extraFileValues,
+            boolean retrain);
 
     native long moduleLoad(
             InputStream is, int[] device, boolean mapLocation, byte[] buffer, long size);
diff --git a/engines/pytorch/pytorch-native/src/main/native/ai_djl_pytorch_jni_PyTorchLibrary_inference.cc b/engines/pytorch/pytorch-native/src/main/native/ai_djl_pytorch_jni_PyTorchLibrary_inference.cc
index c5200081a..11c520a92 100644
--- a/engines/pytorch/pytorch-native/src/main/native/ai_djl_pytorch_jni_PyTorchLibrary_inference.cc
+++ b/engines/pytorch/pytorch-native/src/main/native/ai_djl_pytorch_jni_PyTorchLibrary_inference.cc
@@ -30,9 +30,9 @@ struct JITCallGuard {
 };
 
 JNIEXPORT jlong JNICALL
-Java_ai_djl_pytorch_jni_PyTorchLibrary_moduleLoad__Ljava_lang_String_2_3IZ_3Ljava_lang_String_2_3Ljava_lang_String_2(
+Java_ai_djl_pytorch_jni_PyTorchLibrary_moduleLoad__Ljava_lang_String_2_3IZ_3Ljava_lang_String_2_3Ljava_lang_String_2Z(
     JNIEnv* env, jobject jthis, jstring jpath, jintArray jarray, jboolean jmap_location, jobjectArray jefnames,
-    jobjectArray jefvalues) {
+    jobjectArray jefvalues, jboolean jretrain) {
   API_BEGIN()
   const std::string path = djl::utils::jni::GetStringFromJString(env, jpath);
   const torch::Device device = utils::GetDeviceFromJDevice(env, jarray);
@@ -44,7 +44,25 @@ Java_ai_djl_pytorch_jni_PyTorchLibrary_moduleLoad__Ljava_lang_String_2_3IZ_3Ljav
     map[name] = "";
   }
 
-  JITCallGuard guard;
+  if (!jretrain) {
+    JITCallGuard guard;
+    torch::jit::Module module;
+    if (jmap_location) {
+      module = torch::jit::load(path, device, map);
+      module.eval();
+    } else {
+      module = torch::jit::load(path, torch::nullopt, map);
+      module.eval();
+      module.to(device);
+    }
+    const auto* module_ptr = new torch::jit::Module(module);
+    for (size_t i = 0; i < len; ++i) {
+      auto jname = (jstring) env->GetObjectArrayElement(jefnames, i);
+      auto name = djl::utils::jni::GetStringFromJString(env, jname);
+      env->SetObjectArrayElement(jefvalues, i, env->NewStringUTF(map[name].c_str()));
+    }
+    return reinterpret_cast<uintptr_t>(module_ptr);
+  }
   torch::jit::Module module;
   if (jmap_location) {
     module = torch::jit::load(path, device, map);
diff --git a/examples/src/main/java/ai/djl/examples/training/transferlearning/TransferFreshFruit.java b/examples/src/main/java/ai/djl/examples/training/transferlearning/TransferFreshFruit.java
new file mode 100644
index 000000000..acebbc93e
--- /dev/null
+++ b/examples/src/main/java/ai/djl/examples/training/transferlearning/TransferFreshFruit.java
@@ -0,0 +1,201 @@
+/*
+ * Copyright 2022 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+ * with the License. A copy of the License is located at
+ *
+ * http://aws.amazon.com/apache2.0/
+ *
+ * or in the "license" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES
+ * OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
+ * and limitations under the License.
+ */
+package ai.djl.examples.training.transferlearning;
+
+import ai.djl.Model;
+import ai.djl.ModelException;
+import ai.djl.basicdataset.cv.classification.FruitsFreshAndRotten;
+import ai.djl.engine.Engine;
+import ai.djl.metric.Metrics;
+import ai.djl.modality.cv.transform.CenterCrop;
+import ai.djl.modality.cv.transform.Normalize;
+import ai.djl.modality.cv.transform.OneHot;
+import ai.djl.modality.cv.transform.RandomFlipLeftRight;
+import ai.djl.modality.cv.transform.RandomFlipTopBottom;
+import ai.djl.modality.cv.transform.RandomResizedCrop;
+import ai.djl.modality.cv.transform.Resize;
+import ai.djl.modality.cv.transform.ToTensor;
+import ai.djl.ndarray.NDArray;
+import ai.djl.ndarray.NDList;
+import ai.djl.ndarray.types.Shape;
+import ai.djl.nn.Block;
+import ai.djl.nn.Parameter;
+import ai.djl.nn.SequentialBlock;
+import ai.djl.nn.core.Linear;
+import ai.djl.repository.zoo.Criteria;
+import ai.djl.repository.zoo.ZooModel;
+import ai.djl.training.DefaultTrainingConfig;
+import ai.djl.training.EasyTrain;
+import ai.djl.training.Trainer;
+import ai.djl.training.TrainingResult;
+import ai.djl.training.dataset.Dataset;
+import ai.djl.training.dataset.RandomAccessDataset;
+import ai.djl.training.evaluator.Accuracy;
+import ai.djl.training.listener.SaveModelTrainingListener;
+import ai.djl.training.listener.TrainingListener;
+import ai.djl.training.loss.Loss;
+import ai.djl.training.optimizer.Adam;
+import ai.djl.training.optimizer.Optimizer;
+import ai.djl.training.tracker.FixedPerVarTracker;
+import ai.djl.training.util.ProgressBar;
+import ai.djl.translate.TranslateException;
+import ai.djl.util.Pair;
+
+import java.io.IOException;
+import java.net.URISyntaxException;
+
+public final class TransferFreshFruit {
+
+    private TransferFreshFruit() {}
+
+    public static void main(String[] args)
+            throws IOException, TranslateException, ModelException, URISyntaxException {
+        TransferFreshFruit.runExample(args);
+    }
+
+    public static TrainingResult runExample(String[] args)
+            throws IOException, TranslateException, ModelException, URISyntaxException {
+        boolean retrain = args.length == 1 && "-p".equals(args[0]);
+        Criteria<NDList, NDList> criteria =
+                Criteria.builder()
+                        .setTypes(NDList.class, NDList.class)
+                        .optModelUrls("djl://ai.djl.pytorch/resnet18_embedding")
+                        .optEngine("PyTorch")
+                        .optProgress(new ProgressBar())
+                        .optOption("retrain", String.valueOf(retrain))
+                        .build();
+
+        ZooModel<NDList, NDList> embedding = criteria.loadModel();
+
+        Block baseBlock = embedding.getBlock();
+        Block blocks =
+                new SequentialBlock()
+                        .add(baseBlock)
+                        .addSingleton(nd -> nd.squeeze(new int[] {2, 3}))
+                        .add(Linear.builder().setUnits(2).build()) // linear on which dim?
+                        .addSingleton(nd -> nd.softmax(1));
+
+        Model model = Model.newInstance("TransferFreshFruit");
+        model.setBlock(blocks);
+
+        // Configure trainer
+        DefaultTrainingConfig config = setupTrainingConfig();
+
+        float lr = 0.001f; // Customized learning rate
+        FixedPerVarTracker.Builder learningRateTrackerBuilder =
+                FixedPerVarTracker.builder().setDefaultValue(lr);
+        for (Pair<String, Parameter> paramPair : baseBlock.getParameters()) {
+            learningRateTrackerBuilder.put(paramPair.getValue().getId(), 0.1f * lr);
+        }
+        Optimizer optimizer =
+                Adam.builder().optLearningRateTracker(learningRateTrackerBuilder.build()).build();
+        config.optOptimizer(optimizer);
+
+        Trainer trainer = model.newTrainer(config);
+        trainer.setMetrics(new Metrics());
+
+        // Initialize the parameter shape and value
+        int batchSize = 32;
+        Shape inputShape = new Shape(batchSize, 3, 224, 224);
+        trainer.initialize(inputShape);
+
+        // Data
+        RandomAccessDataset datasetTrain = getData(Dataset.Usage.TRAIN, batchSize);
+        RandomAccessDataset datasetTest = getData(Dataset.Usage.TEST, batchSize);
+
+        // Train
+        EasyTrain.fit(trainer, 10, datasetTrain, datasetTest);
+
+        // Save model
+        // model.save("your-model-path");
+
+        model.close();
+        embedding.close();
+        return trainer.getTrainingResult();
+    }
+
+    private static RandomAccessDataset getData(Dataset.Usage usage, int batchSize)
+            throws TranslateException, IOException {
+        // The dataset is accessible from:
+        // https://www.kaggle.com/datasets/sriramr/fruits-fresh-and-rotten-for-classification
+        float[] mean = {0.485f, 0.456f, 0.406f};
+        float[] std = {0.229f, 0.224f, 0.225f};
+
+        // If the user wants to use local repository, then the dataset can be loaded as follows
+        // Repository repository = Repository.newInstance("banana", Paths.get(LOCAL_FOLDER/{train OR
+        // test}));
+        // FruitsFreshAndRotten dataset =
+        //        FruitsFreshAndRotten.builder()
+        //                            .optRepository(repository)
+        //                            .build()
+        FruitsFreshAndRotten dataset =
+                FruitsFreshAndRotten.builder()
+                        .optUsage(usage)
+                        .optArtifactId("fruit-unittest")
+                        .addTransform(new RandomResizedCrop(256, 256)) // only in training
+                        .addTransform(new RandomFlipTopBottom()) // only in training
+                        .addTransform(new RandomFlipLeftRight()) // only in training
+                        .addTransform(new Resize(256, 256))
+                        .addTransform(new CenterCrop(224, 224))
+                        .addTransform(new ToTensor())
+                        .addTransform(new Normalize(mean, std))
+                        .addTargetTransform(new OneHot(2))
+                        .setSampling(batchSize, true)
+                        .build();
+
+        dataset.prepare();
+        return dataset;
+    }
+
+    private static DefaultTrainingConfig setupTrainingConfig() {
+        String outputDir = "build/fruits";
+        SaveModelTrainingListener listener = new SaveModelTrainingListener(outputDir);
+        listener.setSaveModelCallback(
+                trainer -> {
+                    TrainingResult result = trainer.getTrainingResult();
+                    Model model = trainer.getModel();
+                    float accuracy = result.getValidateEvaluation("Accuracy");
+                    model.setProperty("Accuracy", String.format("%.5f", accuracy));
+                    model.setProperty("Loss", String.format("%.5f", result.getValidateLoss()));
+                });
+
+        return new DefaultTrainingConfig(new SoftmaxCrossEntropy("SoftmaxCrossEntropy"))
+                .addEvaluator(new Accuracy())
+                .optDevices(Engine.getInstance().getDevices(1))
+                .addTrainingListeners(TrainingListener.Defaults.logging(outputDir))
+                .addTrainingListeners(listener);
+    }
+
+    private static class SoftmaxCrossEntropy extends Loss {
+
+        /**
+         * Base class for metric with abstract update methods.
+         *
+         * @param name The display name of the Loss
+         */
+        public SoftmaxCrossEntropy(String name) {
+            super(name);
+        }
+
+        /** {@inheritDoc} */
+        @Override
+        public NDArray evaluate(NDList labels, NDList predictions) {
+            // Here the labels are supposed to be one-hot
+            int classAxis = -1;
+            NDArray pred = predictions.singletonOrThrow().log();
+            NDArray lab = labels.singletonOrThrow().reshape(pred.getShape());
+            NDArray loss = pred.mul(lab).neg().sum(new int[] {classAxis}, true);
+            return loss.mean();
+        }
+    }
+}


Test patch:
diff --git a/api/src/test/java/ai/djl/training/LearningRateTest.java b/api/src/test/java/ai/djl/training/LearningRateTest.java
index dc0cd099c..12a8a1581 100644
--- a/api/src/test/java/ai/djl/training/LearningRateTest.java
+++ b/api/src/test/java/ai/djl/training/LearningRateTest.java
@@ -16,6 +16,7 @@ package ai.djl.training;
 import ai.djl.training.tracker.CosineTracker;
 import ai.djl.training.tracker.CyclicalTracker;
 import ai.djl.training.tracker.FactorTracker;
+import ai.djl.training.tracker.FixedPerVarTracker;
 import ai.djl.training.tracker.MultiFactorTracker;
 import ai.djl.training.tracker.Tracker;
 
@@ -143,6 +144,22 @@ public class LearningRateTest {
         Assert.assertEquals(cyclicalTrackerCustom.getNewValue(500), 0.04f, epsilon);
     }
 
+    @Test
+    public void testFixedPerVarTracker() {
+        float lr = 0.001f; // Customized learning rate
+        FixedPerVarTracker.Builder fixedPerVarTrackerBuilder =
+                FixedPerVarTracker.builder().setDefaultValue(lr);
+        String[] parameterIds = {"id1", "id2"};
+        for (String param : parameterIds) {
+            fixedPerVarTrackerBuilder.put(param, 0.1f * lr);
+        }
+        FixedPerVarTracker fixedPerVarTracker = fixedPerVarTrackerBuilder.build();
+
+        Assert.assertEquals(fixedPerVarTracker.getNewValue("id1", 0), 0.1f * lr);
+        Assert.assertEquals(fixedPerVarTracker.getNewValue("id2", 1), 0.1f * lr);
+        Assert.assertEquals(fixedPerVarTracker.getNewValue("unknown", 2), lr);
+    }
+
     private static class CustomScaleFunction implements CyclicalTracker.ScaleFunction {
         @Override
         public float func(int steps) {
diff --git a/basicdataset/src/test/java/ai/djl/basicdataset/FruitsFreshAndRottenTest.java b/basicdataset/src/test/java/ai/djl/basicdataset/FruitsFreshAndRottenTest.java
new file mode 100644
index 000000000..14a4c1ede
--- /dev/null
+++ b/basicdataset/src/test/java/ai/djl/basicdataset/FruitsFreshAndRottenTest.java
@@ -0,0 +1,58 @@
+/*
+ * Copyright 2022 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+ * with the License. A copy of the License is located at
+ *
+ * http://aws.amazon.com/apache2.0/
+ *
+ * or in the "license" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES
+ * OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
+ * and limitations under the License.
+ */
+package ai.djl.basicdataset;
+
+import ai.djl.Model;
+import ai.djl.basicdataset.cv.classification.FruitsFreshAndRotten;
+import ai.djl.modality.cv.transform.Resize;
+import ai.djl.nn.Blocks;
+import ai.djl.training.DefaultTrainingConfig;
+import ai.djl.training.Trainer;
+import ai.djl.training.TrainingConfig;
+import ai.djl.training.dataset.Batch;
+import ai.djl.training.dataset.Dataset;
+import ai.djl.training.loss.Loss;
+import ai.djl.translate.TranslateException;
+
+import org.testng.Assert;
+import org.testng.annotations.Test;
+
+import java.io.IOException;
+
+public class FruitsFreshAndRottenTest {
+
+    @Test
+    public void testFruitsFreshAndRottenRemote() throws IOException, TranslateException {
+        TrainingConfig config = new DefaultTrainingConfig(Loss.softmaxCrossEntropyLoss());
+
+        try (Model model = Model.newInstance("model")) {
+            model.setBlock(Blocks.identityBlock());
+
+            FruitsFreshAndRotten dataset =
+                    FruitsFreshAndRotten.builder()
+                            .optUsage(Dataset.Usage.TEST)
+                            .optArtifactId("fruit-unittest")
+                            .addTransform(new Resize(256, 256))
+                            .setSampling(32, true)
+                            .build();
+            dataset.prepare();
+
+            try (Trainer trainer = model.newTrainer(config)) {
+                Batch batch = trainer.iterateDataset(dataset).iterator().next();
+                Assert.assertEquals(batch.getData().size(), 1);
+                Assert.assertEquals(batch.getLabels().size(), 1);
+                batch.close();
+            }
+        }
+    }
+}
diff --git a/engines/mxnet/mxnet-engine/src/test/java/ai/djl/mxnet/integration/MxParameterStoreTest.java b/engines/mxnet/mxnet-engine/src/test/java/ai/djl/mxnet/integration/MxParameterStoreTest.java
index 65ed1581b..5c1005eea 100644
--- a/engines/mxnet/mxnet-engine/src/test/java/ai/djl/mxnet/integration/MxParameterStoreTest.java
+++ b/engines/mxnet/mxnet-engine/src/test/java/ai/djl/mxnet/integration/MxParameterStoreTest.java
@@ -23,6 +23,7 @@ import ai.djl.testing.Assertions;
 import ai.djl.testing.TestRequirements;
 import ai.djl.training.ParameterServer;
 import ai.djl.training.optimizer.Optimizer;
+import ai.djl.training.tracker.ParameterTracker;
 import ai.djl.training.tracker.Tracker;
 
 import org.testng.Assert;
@@ -99,7 +100,7 @@ public class MxParameterStoreTest {
 
     private static class TestOptimizer extends Optimizer {
 
-        private Tracker learningRateTracker;
+        private ParameterTracker learningRateTracker;
         int updateCount;
 
         protected TestOptimizer(TestOptimizer.Builder builder) {
@@ -111,7 +112,7 @@ public class MxParameterStoreTest {
         @Override
         public void update(String parameterId, NDArray weight, NDArray grad) {
             weight.addi(
-                    grad.mul(learningRateTracker.getNewValue(0))
+                    grad.mul(learningRateTracker.getNewValue(parameterId, 0))
                             .toDevice(weight.getDevice(), false));
             updateCount++;
         }
diff --git a/examples/src/test/java/ai/djl/examples/training/TransferFreshFruitTest.java b/examples/src/test/java/ai/djl/examples/training/TransferFreshFruitTest.java
new file mode 100644
index 000000000..39f148a20
--- /dev/null
+++ b/examples/src/test/java/ai/djl/examples/training/TransferFreshFruitTest.java
@@ -0,0 +1,43 @@
+/*
+ * Copyright 2022 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+ * with the License. A copy of the License is located at
+ *
+ * http://aws.amazon.com/apache2.0/
+ *
+ * or in the "license" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES
+ * OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
+ * and limitations under the License.
+ */
+package ai.djl.examples.training;
+
+import ai.djl.ModelException;
+import ai.djl.examples.training.transferlearning.TransferFreshFruit;
+import ai.djl.testing.TestRequirements;
+import ai.djl.training.TrainingResult;
+import ai.djl.translate.TranslateException;
+
+import org.testng.Assert;
+import org.testng.annotations.Test;
+
+import java.io.IOException;
+import java.net.URISyntaxException;
+
+public class TransferFreshFruitTest {
+
+    @Test
+    public void testTransferFreshFruit()
+            throws ModelException, TranslateException, IOException, URISyntaxException {
+        TestRequirements.engine("PyTorch");
+
+        String[][] args = {{}, {"-p"}};
+        TrainingResult result;
+        for (String[] arg : args) {
+            result = TransferFreshFruit.runExample(arg);
+            Assert.assertNotNull(result);
+            Assert.assertTrue(result.getEvaluations().get("validate_Accuracy") > 0.9f);
+            Assert.assertTrue(result.getEvaluations().get("train_Accuracy") > 0.9f);
+        }
+    }
+}