Repository: deepjavalibrary/djl
Pull request #2303 (2023-01-07 01:26:11)

[api] Refactor Engine class
## 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


codecov-commenter: # [Codecov](https://codecov.io/gh/deepjavalibrary/djl/pull/2303?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: **74.24**% // Increases project coverage by **`+2.15%`** :tada:
> Coverage data is based on head [(`19a6d32`)](https://codecov.io/gh/deepjavalibrary/djl/pull/2303?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: 74.97% of modified lines in pull request are covered.

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


```diff
@@             Coverage Diff              @@
##             master    #2303      +/-   ##
============================================
+ Coverage     72.08%   74.24%   +2.15%     
- Complexity     5126     6680    +1554     
============================================
  Files           473      656     +183     
  Lines         21970    29060    +7090     
  Branches       2351     3002     +651     
============================================
+ Hits          15838    21576    +5738     
- Misses         4925     6019    +1094     
- Partials       1207     1465     +258     
```


| [Impacted Files](https://codecov.io/gh/deepjavalibrary/djl/pull/2303?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/2303?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/2303?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/2303?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/2303?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/2303?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/2303?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: |
| [.../modality/cv/translator/ImageFeatureExtractor.java](https://codecov.io/gh/deepjavalibrary/djl/pull/2303?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3RyYW5zbGF0b3IvSW1hZ2VGZWF0dXJlRXh0cmFjdG9yLmphdmE=) | `0.00% <0.00%> (ø)` | |
| [.../ai/djl/modality/cv/translator/YoloTranslator.java](https://codecov.io/gh/deepjavalibrary/djl/pull/2303?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: |
| [...ain/java/ai/djl/modality/cv/util/NDImageUtils.java](https://codecov.io/gh/deepjavalibrary/djl/pull/2303?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3V0aWwvTkRJbWFnZVV0aWxzLmphdmE=) | `59.21% <0.00%> (ø)` | |
| [api/src/main/java/ai/djl/modality/nlp/Decoder.java](https://codecov.io/gh/deepjavalibrary/djl/pull/2303?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L25scC9EZWNvZGVyLmphdmE=) | `63.63% <ø> (ø)` | |
| ... and [599 more](https://codecov.io/gh/deepjavalibrary/djl/pull/2303?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/2303?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/engine/Engine.java b/api/src/main/java/ai/djl/engine/Engine.java
index 73d19779b..3b29a9295 100644
--- a/api/src/main/java/ai/djl/engine/Engine.java
+++ b/api/src/main/java/ai/djl/engine/Engine.java
@@ -21,6 +21,7 @@ import ai.djl.training.LocalParameterServer;
 import ai.djl.training.ParameterServer;
 import ai.djl.training.optimizer.Optimizer;
 import ai.djl.util.Ec2Utils;
+import ai.djl.util.RandomUtils;
 import ai.djl.util.Utils;
 import ai.djl.util.cuda.CudaUtils;
 
@@ -271,7 +272,9 @@ public abstract class Engine {
      * @param manager the manager to manage parameters
      * @return Empty {@link SymbolBlock} for static graph
      */
-    public abstract SymbolBlock newSymbolBlock(NDManager manager);
+    public SymbolBlock newSymbolBlock(NDManager manager) {
+        throw new UnsupportedOperationException("Not supported.");
+    }
 
     /**
      * Constructs a new model.
@@ -304,7 +307,9 @@ public abstract class Engine {
      *
      * @return a new instance of {@link GradientCollector}
      */
-    public abstract GradientCollector newGradientCollector();
+    public GradientCollector newGradientCollector() {
+        throw new UnsupportedOperationException("Not supported.");
+    }
 
     /**
      * Returns a new instance of {@link ParameterServer}.
@@ -326,6 +331,7 @@ public abstract class Engine {
      */
     public void setRandomSeed(int seed) {
         this.seed = seed;
+        RandomUtils.RANDOM.setSeed(seed);
     }
 
     /**
diff --git a/api/src/main/java/ai/djl/ndarray/BaseNDManager.java b/api/src/main/java/ai/djl/ndarray/BaseNDManager.java
index 59b69fbe7..abbc3c316 100644
--- a/api/src/main/java/ai/djl/ndarray/BaseNDManager.java
+++ b/api/src/main/java/ai/djl/ndarray/BaseNDManager.java
@@ -32,12 +32,9 @@ import java.nio.LongBuffer;
 import java.nio.ShortBuffer;
 import java.nio.charset.Charset;
 import java.nio.file.Path;
-import java.util.List;
 import java.util.UUID;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.atomic.AtomicBoolean;
-import java.util.stream.Collectors;
-import java.util.stream.Stream;
 
 /** {@code BaseNDManager} is the default implementation of {@link NDManager}. */
 public abstract class BaseNDManager implements NDManager {
@@ -307,30 +304,6 @@ public abstract class BaseNDManager implements NDManager {
         return device;
     }
 
-    /** {@inheritDoc} */
-    @Override
-    public List<NDArray> getManagedArrays() {
-        return Stream.concat(
-                        // Main resources
-                        resources.values().stream()
-                                .flatMap(
-                                        r -> {
-                                            if (r instanceof NDResource) {
-                                                return ((NDResource) r)
-                                                        .getResourceNDArrays().stream();
-                                            } else if (r instanceof NDManager) {
-                                                return ((NDManager) r).getManagedArrays().stream();
-                                            } else {
-                                                return Stream.empty();
-                                            }
-                                        }),
-
-                        // Temp resouces
-                        tempResources.values().stream()
-                                .flatMap(tr -> tr.resource.getResourceNDArrays().stream()))
-                .collect(Collectors.toList());
-    }
-
     /** {@inheritDoc} */
     @Override
     public String toString() {
@@ -348,6 +321,9 @@ public abstract class BaseNDManager implements NDManager {
     /** {@inheritDoc} */
     @Override
     public synchronized void attachInternal(String resourceId, AutoCloseable resource) {
+        if (this instanceof SystemNDManager) {
+            return;
+        }
         if (capped.get()) {
             throw new IllegalStateException("NDManager is capped for addition of resources.");
         }
@@ -357,6 +333,9 @@ public abstract class BaseNDManager implements NDManager {
     /** {@inheritDoc} */
     @Override
     public synchronized void attachUncappedInternal(String resourceId, AutoCloseable resource) {
+        if (this instanceof SystemNDManager) {
+            return;
+        }
         if (closed.get()) {
             throw new IllegalStateException("NDManager has been closed already.");
         }
@@ -383,8 +362,7 @@ public abstract class BaseNDManager implements NDManager {
     public void tempAttachInternal(
             NDManager originalManager, String resourceId, NDResource resource) {
         if (this instanceof SystemNDManager) {
-            throw new IllegalStateException(
-                    "System manager cannot be temp attached because it can't be closed..");
+            return;
         }
         if (closed.get()) {
             throw new IllegalStateException("NDManager has been closed already.");
@@ -395,6 +373,9 @@ public abstract class BaseNDManager implements NDManager {
     /** {@inheritDoc} */
     @Override
     public synchronized void detachInternal(String resourceId) {
+        if (this instanceof SystemNDManager) {
+            return;
+        }
         if (closed.get()) {
             // This may happen in the middle of BaseNDManager.close()
             return;
@@ -421,13 +402,26 @@ public abstract class BaseNDManager implements NDManager {
         throw new UnsupportedOperationException("Not supported!");
     }
 
+    /** {@inheritDoc} */
+    @Override
+    public void zeroGradients() {
+        for (AutoCloseable res : resources.values()) {
+            if (res instanceof NDManager) {
+                ((NDManager) res).zeroGradients();
+            } else if (res instanceof NDArray) {
+                NDArray array = (NDArray) res;
+                if (array.hasGradient()) {
+                    array.getGradient().subi(array.getGradient());
+                }
+            }
+        }
+    }
+
     /** {@inheritDoc} */
     @Override
     public void close() {
         if (this instanceof SystemNDManager) {
-            throw new IllegalStateException(
-                    "The SystemNDManager can not be closed. It is global and lives for the duration"
-                            + " of the process");
+            return;
         }
         if (!closed.getAndSet(true)) {
             for (AutoCloseable closeable : resources.values()) {
diff --git a/api/src/main/java/ai/djl/ndarray/NDArray.java b/api/src/main/java/ai/djl/ndarray/NDArray.java
index ee67c7f78..1adfd59ac 100644
--- a/api/src/main/java/ai/djl/ndarray/NDArray.java
+++ b/api/src/main/java/ai/djl/ndarray/NDArray.java
@@ -30,8 +30,6 @@ import java.nio.LongBuffer;
 import java.nio.charset.Charset;
 import java.nio.charset.StandardCharsets;
 import java.util.Arrays;
-import java.util.Collections;
-import java.util.List;
 import java.util.function.Function;
 import java.util.stream.IntStream;
 import java.util.stream.LongStream;
@@ -4687,12 +4685,6 @@ public interface NDArray extends NDResource, BytesSupplier {
      */
     NDArray erfinv();
 
-    /** {@inheritDoc} */
-    @Override
-    default List<NDArray> getResourceNDArrays() {
-        return Collections.singletonList(this);
-    }
-
     /**
      * Returns an internal representative of Native {@code NDArray}.
      *
diff --git a/api/src/main/java/ai/djl/ndarray/NDList.java b/api/src/main/java/ai/djl/ndarray/NDList.java
index 36c73bb42..c88c93337 100644
--- a/api/src/main/java/ai/djl/ndarray/NDList.java
+++ b/api/src/main/java/ai/djl/ndarray/NDList.java
@@ -27,7 +27,6 @@ import java.nio.ByteBuffer;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collection;
-import java.util.List;
 import java.util.zip.ZipEntry;
 import java.util.zip.ZipInputStream;
 import java.util.zip.ZipOutputStream;
@@ -270,12 +269,6 @@ public class NDList extends ArrayList<NDArray> implements NDResource, BytesSuppl
         return head().getManager();
     }
 
-    /** {@inheritDoc} */
-    @Override
-    public List<NDArray> getResourceNDArrays() {
-        return this;
-    }
-
     /** {@inheritDoc} */
     @Override
     public void attach(NDManager manager) {
diff --git a/api/src/main/java/ai/djl/ndarray/NDManager.java b/api/src/main/java/ai/djl/ndarray/NDManager.java
index 1d0e869cf..d61ba4387 100644
--- a/api/src/main/java/ai/djl/ndarray/NDManager.java
+++ b/api/src/main/java/ai/djl/ndarray/NDManager.java
@@ -34,7 +34,6 @@ import java.nio.LongBuffer;
 import java.nio.charset.Charset;
 import java.nio.charset.StandardCharsets;
 import java.nio.file.Path;
-import java.util.List;
 
 /**
  * NDArray managers are used to create <I>NDArrays</I> (n-dimensional array on native engine).
@@ -1535,13 +1534,6 @@ public interface NDManager extends AutoCloseable {
      */
     Device getDevice();
 
-    /**
-     * Returns all {@link NDArray}s managed by this manager (including recursively).
-     *
-     * @return all {@link NDArray}s managed by this manager (including recursively)
-     */
-    List<NDArray> getManagedArrays();
-
     /**
      * Attaches a resource to this {@code NDManager}.
      *
@@ -1676,6 +1668,9 @@ public interface NDManager extends AutoCloseable {
      */
     Engine getEngine();
 
+    /** Sets all the gradients within the NDManager to zero. */
+    void zeroGradients();
+
     /** {@inheritDoc} */
     @Override
     void close();
diff --git a/api/src/main/java/ai/djl/ndarray/NDResource.java b/api/src/main/java/ai/djl/ndarray/NDResource.java
index 02d09bb77..72d709c8d 100644
--- a/api/src/main/java/ai/djl/ndarray/NDResource.java
+++ b/api/src/main/java/ai/djl/ndarray/NDResource.java
@@ -12,8 +12,6 @@
  */
 package ai.djl.ndarray;
 
-import java.util.List;
-
 /** An object which is managed by an {@link NDManager} and tracks the manager it is attached to. */
 public interface NDResource extends AutoCloseable {
 
@@ -24,13 +22,6 @@ public interface NDResource extends AutoCloseable {
      */
     NDManager getManager();
 
-    /**
-     * Returns the {@link NDArray} or {@link NDArray}s contained within this resource.
-     *
-     * @return the {@link NDArray} or {@link NDArray}s contained within this resource
-     */
-    List<NDArray> getResourceNDArrays();
-
     /**
      * Attaches this {@link NDResource} to the specified {@link NDManager}.
      *
diff --git a/api/src/main/java/ai/djl/training/GradientCollector.java b/api/src/main/java/ai/djl/training/GradientCollector.java
index 0e16e94d8..e1a3cd99c 100644
--- a/api/src/main/java/ai/djl/training/GradientCollector.java
+++ b/api/src/main/java/ai/djl/training/GradientCollector.java
@@ -21,14 +21,6 @@ import ai.djl.ndarray.NDArray;
  * performed within the try-with-resources are recorded and the variables marked. When {@link
  * #backward(NDArray) backward function} is called, gradients are collected w.r.t previously marked
  * variables.
- *
- * <p>The typical behavior is to open up a gradient collector during each batch and close it during
- * the end of the batch. In this way, the gradient is reset between batches. If the gradient
- * collector is left open for multiple calls to backwards, the gradients collected are accumulated
- * and added together.
- *
- * <p>Due to limitations in most engines, the gradient collectors are global. This means that only
- * one can be used at a time. If multiple are opened, an error will be thrown.
  */
 public interface GradientCollector extends AutoCloseable {
 
@@ -39,9 +31,6 @@ public interface GradientCollector extends AutoCloseable {
      */
     void backward(NDArray target);
 
-    /** Sets all the gradients within the engine to zero. */
-    void zeroGradients();
-
     /** {@inheritDoc} */
     @Override
     void close();
diff --git a/api/src/main/java/ai/djl/util/passthrough/PassthroughNDManager.java b/api/src/main/java/ai/djl/util/passthrough/PassthroughNDManager.java
index bde8a8913..7b35184af 100644
--- a/api/src/main/java/ai/djl/util/passthrough/PassthroughNDManager.java
+++ b/api/src/main/java/ai/djl/util/passthrough/PassthroughNDManager.java
@@ -26,8 +26,6 @@ import java.nio.Buffer;
 import java.nio.ByteBuffer;
 import java.nio.charset.Charset;
 import java.nio.file.Path;
-import java.util.Collections;
-import java.util.List;
 
 /** An {@link NDManager} that does nothing, for use in extensions and hybrid engines. */
 public final class PassthroughNDManager implements NDManager {
@@ -249,12 +247,6 @@ public final class PassthroughNDManager implements NDManager {
         return Device.cpu();
     }
 
-    /** {@inheritDoc} */
-    @Override
-    public List<NDArray> getManagedArrays() {
-        return Collections.emptyList();
-    }
-
     /** {@inheritDoc} */
     @Override
     public void attachInternal(String resourceId, AutoCloseable resource) {
@@ -299,6 +291,12 @@ public final class PassthroughNDManager implements NDManager {
         return null;
     }
 
+    /** {@inheritDoc} */
+    @Override
+    public void zeroGradients() {
+        throw new UnsupportedOperationException(UNSUPPORTED);
+    }
+
     /** {@inheritDoc} */
     @Override
     public void close() {}
diff --git a/engines/dlr/dlr-engine/src/main/java/ai/djl/dlr/engine/DlrEngine.java b/engines/dlr/dlr-engine/src/main/java/ai/djl/dlr/engine/DlrEngine.java
index ee8bf0d4f..727bed9cf 100644
--- a/engines/dlr/dlr-engine/src/main/java/ai/djl/dlr/engine/DlrEngine.java
+++ b/engines/dlr/dlr-engine/src/main/java/ai/djl/dlr/engine/DlrEngine.java
@@ -19,8 +19,6 @@ import ai.djl.dlr.jni.LibUtils;
 import ai.djl.engine.Engine;
 import ai.djl.engine.EngineException;
 import ai.djl.ndarray.NDManager;
-import ai.djl.nn.SymbolBlock;
-import ai.djl.training.GradientCollector;
 
 /**
  * The {@code DlrEngine} is an implementation of the {@link Engine} based on the <a
@@ -86,12 +84,6 @@ public final class DlrEngine extends Engine {
         return false;
     }
 
-    /** {@inheritDoc} */
-    @Override
-    public SymbolBlock newSymbolBlock(NDManager manager) {
-        throw new UnsupportedOperationException("DLR does not support empty SymbolBlock");
-    }
-
     /** {@inheritDoc} */
     @Override
     public Model newModel(String name, Device device) {
@@ -113,16 +105,4 @@ public final class DlrEngine extends Engine {
     public NDManager newBaseManager(Device device) {
         return DlrNDManager.getSystemManager().newSubManager(device);
     }
-
-    /** {@inheritDoc} */
-    @Override
-    public GradientCollector newGradientCollector() {
-        throw new UnsupportedOperationException("Not supported for DLR");
-    }
-
-    /** {@inheritDoc} */
-    @Override
-    public void setRandomSeed(int seed) {
-        throw new UnsupportedOperationException("Not supported for DLR");
-    }
 }
diff --git a/engines/ml/lightgbm/src/main/java/ai/djl/ml/lightgbm/LgbmEngine.java b/engines/ml/lightgbm/src/main/java/ai/djl/ml/lightgbm/LgbmEngine.java
index 23d0c1adb..3f30e8b7f 100644
--- a/engines/ml/lightgbm/src/main/java/ai/djl/ml/lightgbm/LgbmEngine.java
+++ b/engines/ml/lightgbm/src/main/java/ai/djl/ml/lightgbm/LgbmEngine.java
@@ -18,8 +18,6 @@ import ai.djl.engine.Engine;
 import ai.djl.engine.EngineException;
 import ai.djl.ml.lightgbm.jni.LibUtils;
 import ai.djl.ndarray.NDManager;
-import ai.djl.nn.SymbolBlock;
-import ai.djl.training.GradientCollector;
 
 import java.io.IOException;
 
@@ -89,12 +87,6 @@ public final class LgbmEngine extends Engine {
         return false;
     }
 
-    /** {@inheritDoc} */
-    @Override
-    public SymbolBlock newSymbolBlock(NDManager manager) {
-        throw new UnsupportedOperationException("LightGBM does not support empty symbol block");
-    }
-
     /** {@inheritDoc} */
     @Override
     public Model newModel(String name, Device device) {
@@ -112,16 +104,4 @@ public final class LgbmEngine extends Engine {
     public NDManager newBaseManager(Device device) {
         return LgbmNDManager.getSystemManager().newSubManager(device);
     }
-
-    /** {@inheritDoc} */
-    @Override
-    public GradientCollector newGradientCollector() {
-        throw new UnsupportedOperationException("Not supported for LightGBM");
-    }
-
-    /** {@inheritDoc} */
-    @Override
-    public void setRandomSeed(int seed) {
-        throw new UnsupportedOperationException("Not supported for LightGBM");
-    }
 }
diff --git a/engines/ml/xgboost/src/main/java/ai/djl/ml/xgboost/XgbEngine.java b/engines/ml/xgboost/src/main/java/ai/djl/ml/xgboost/XgbEngine.java
index 6449ee5f3..70e499ebb 100644
--- a/engines/ml/xgboost/src/main/java/ai/djl/ml/xgboost/XgbEngine.java
+++ b/engines/ml/xgboost/src/main/java/ai/djl/ml/xgboost/XgbEngine.java
@@ -17,8 +17,6 @@ import ai.djl.Model;
 import ai.djl.engine.Engine;
 import ai.djl.engine.StandardCapabilities;
 import ai.djl.ndarray.NDManager;
-import ai.djl.nn.SymbolBlock;
-import ai.djl.training.GradientCollector;
 
 import ml.dmlc.xgboost4j.java.JniUtils;
 
@@ -101,12 +99,6 @@ public final class XgbEngine extends Engine {
         return false;
     }
 
-    /** {@inheritDoc} */
-    @Override
-    public SymbolBlock newSymbolBlock(NDManager manager) {
-        throw new UnsupportedOperationException("XGBoost does not support empty symbol block");
-    }
-
     /** {@inheritDoc} */
     @Override
     public Model newModel(String name, Device device) {
@@ -125,18 +117,6 @@ public final class XgbEngine extends Engine {
         return XgbNDManager.getSystemManager().newSubManager(device);
     }
 
-    /** {@inheritDoc} */
-    @Override
-    public GradientCollector newGradientCollector() {
-        throw new UnsupportedOperationException("Not supported for XGBoost");
-    }
-
-    /** {@inheritDoc} */
-    @Override
-    public void setRandomSeed(int seed) {
-        throw new UnsupportedOperationException("Not supported for XGBoost");
-    }
-
     /** {@inheritDoc} */
     @Override
     public String toString() {
diff --git a/engines/mxnet/mxnet-engine/src/main/java/ai/djl/mxnet/engine/MxEngine.java b/engines/mxnet/mxnet-engine/src/main/java/ai/djl/mxnet/engine/MxEngine.java
index 7f4d39e83..b95ff8517 100644
--- a/engines/mxnet/mxnet-engine/src/main/java/ai/djl/mxnet/engine/MxEngine.java
+++ b/engines/mxnet/mxnet-engine/src/main/java/ai/djl/mxnet/engine/MxEngine.java
@@ -24,7 +24,6 @@ import ai.djl.training.GradientCollector;
 import ai.djl.training.LocalParameterServer;
 import ai.djl.training.ParameterServer;
 import ai.djl.training.optimizer.Optimizer;
-import ai.djl.util.RandomUtils;
 import ai.djl.util.Utils;
 
 import java.io.FileNotFoundException;
@@ -157,7 +156,6 @@ public final class MxEngine extends Engine {
     public void setRandomSeed(int seed) {
         super.setRandomSeed(seed);
         JnaUtils.randomSeed(seed);
-        RandomUtils.RANDOM.setSeed(seed);
     }
 
     /** {@inheritDoc} */
diff --git a/engines/mxnet/mxnet-engine/src/main/java/ai/djl/mxnet/engine/MxGradientCollector.java b/engines/mxnet/mxnet-engine/src/main/java/ai/djl/mxnet/engine/MxGradientCollector.java
index a8a183de3..5e1ad8369 100644
--- a/engines/mxnet/mxnet-engine/src/main/java/ai/djl/mxnet/engine/MxGradientCollector.java
+++ b/engines/mxnet/mxnet-engine/src/main/java/ai/djl/mxnet/engine/MxGradientCollector.java
@@ -19,7 +19,7 @@ import ai.djl.ndarray.NDManager;
 import ai.djl.training.GradientCollector;
 
 /** {@code MxGradientCollector} is the MXNet implementation of {@link GradientCollector}. */
-public final class MxGradientCollector implements GradientCollector {
+public class MxGradientCollector implements GradientCollector {
 
     /**
      * Constructs an {@code MxGradientCollector} and enables training data collection for
@@ -116,15 +116,4 @@ public final class MxGradientCollector implements GradientCollector {
     private void backward(NDArray array, boolean retainGraph) {
         JnaUtils.autogradBackward(new NDList(array), retainGraph ? 1 : 0);
     }
-
-    /** {@inheritDoc} */
-    @Override
-    public void zeroGradients() {
-        NDManager systemManager = MxNDManager.getSystemManager();
-        for (NDArray array : systemManager.getManagedArrays()) {
-            if (array.hasGradient()) {
-                array.getGradient().subi(array.getGradient());
-            }
-        }
-    }
 }
diff --git a/engines/onnxruntime/onnxruntime-engine/src/main/java/ai/djl/onnxruntime/engine/OrtEngine.java b/engines/onnxruntime/onnxruntime-engine/src/main/java/ai/djl/onnxruntime/engine/OrtEngine.java
index 6bf773e15..3df5a7c01 100644
--- a/engines/onnxruntime/onnxruntime-engine/src/main/java/ai/djl/onnxruntime/engine/OrtEngine.java
+++ b/engines/onnxruntime/onnxruntime-engine/src/main/java/ai/djl/onnxruntime/engine/OrtEngine.java
@@ -18,8 +18,6 @@ import ai.djl.Model;
 import ai.djl.engine.Engine;
 import ai.djl.engine.StandardCapabilities;
 import ai.djl.ndarray.NDManager;
-import ai.djl.nn.SymbolBlock;
-import ai.djl.training.GradientCollector;
 import ai.onnxruntime.OrtEnvironment;
 import ai.onnxruntime.OrtException;
 import ai.onnxruntime.OrtLoggingLevel;
@@ -124,12 +122,6 @@ public final class OrtEngine extends Engine {
         return new OrtModel(name, newBaseManager(device), env);
     }
 
-    /** {@inheritDoc} */
-    @Override
-    public SymbolBlock newSymbolBlock(NDManager manager) {
-        throw new UnsupportedOperationException("ONNXRuntime does not support empty SymbolBlock");
-    }
-
     /** {@inheritDoc} */
     @Override
     public NDManager newBaseManager() {
@@ -142,18 +134,6 @@ public final class OrtEngine extends Engine {
         return OrtNDManager.getSystemManager().newSubManager(device);
     }
 
-    /** {@inheritDoc} */
-    @Override
-    public GradientCollector newGradientCollector() {
-        throw new UnsupportedOperationException("Not supported for ONNX Runtime");
-    }
-
-    /** {@inheritDoc} */
-    @Override
-    public void setRandomSeed(int seed) {
-        throw new UnsupportedOperationException("Not supported for ONNX Runtime");
-    }
-
     /** {@inheritDoc} */
     @Override
     public String toString() {
diff --git a/engines/paddlepaddle/paddlepaddle-engine/src/main/java/ai/djl/paddlepaddle/engine/PpEngine.java b/engines/paddlepaddle/paddlepaddle-engine/src/main/java/ai/djl/paddlepaddle/engine/PpEngine.java
index e48249159..06ac07c6d 100644
--- a/engines/paddlepaddle/paddlepaddle-engine/src/main/java/ai/djl/paddlepaddle/engine/PpEngine.java
+++ b/engines/paddlepaddle/paddlepaddle-engine/src/main/java/ai/djl/paddlepaddle/engine/PpEngine.java
@@ -16,9 +16,7 @@ import ai.djl.Device;
 import ai.djl.Model;
 import ai.djl.engine.Engine;
 import ai.djl.ndarray.NDManager;
-import ai.djl.nn.SymbolBlock;
 import ai.djl.paddlepaddle.jni.LibUtils;
-import ai.djl.training.GradientCollector;
 
 import java.io.IOException;
 import java.io.InputStream;
@@ -98,12 +96,6 @@ public final class PpEngine extends Engine {
         return new PpModel(name, device, newBaseManager(device));
     }
 
-    /** {@inheritDoc} */
-    @Override
-    public SymbolBlock newSymbolBlock(NDManager manager) {
-        throw new UnsupportedOperationException("PaddlePaddle does not support empty SymbolBlock");
-    }
-
     /** {@inheritDoc} */
     @Override
     public NDManager newBaseManager() {
@@ -115,16 +107,4 @@ public final class PpEngine extends Engine {
     public NDManager newBaseManager(Device device) {
         return PpNDManager.getSystemManager().newSubManager(device);
     }
-
-    /** {@inheritDoc} */
-    @Override
-    public GradientCollector newGradientCollector() {
-        throw new UnsupportedOperationException("Not supported for PaddlePaddle");
-    }
-
-    /** {@inheritDoc} */
-    @Override
-    public void setRandomSeed(int seed) {
-        throw new UnsupportedOperationException("Not supported for PaddlePaddle");
-    }
 }
diff --git a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtEngine.java b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtEngine.java
index d0f7a1013..3add5a145 100644
--- a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtEngine.java
+++ b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtEngine.java
@@ -21,7 +21,6 @@ import ai.djl.nn.SymbolBlock;
 import ai.djl.pytorch.jni.JniUtils;
 import ai.djl.pytorch.jni.LibUtils;
 import ai.djl.training.GradientCollector;
-import ai.djl.util.RandomUtils;
 import ai.djl.util.Utils;
 
 import org.slf4j.Logger;
@@ -151,7 +150,6 @@ public final class PtEngine extends Engine {
     public void setRandomSeed(int seed) {
         super.setRandomSeed(seed);
         JniUtils.setSeed(seed);
-        RandomUtils.RANDOM.setSeed(seed);
     }
 
     /** {@inheritDoc} */
diff --git a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtGradientCollector.java b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtGradientCollector.java
index 5ef573cee..fad7755fa 100644
--- a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtGradientCollector.java
+++ b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtGradientCollector.java
@@ -14,31 +14,18 @@
 package ai.djl.pytorch.engine;
 
 import ai.djl.ndarray.NDArray;
-import ai.djl.ndarray.NDManager;
 import ai.djl.pytorch.jni.JniUtils;
 import ai.djl.training.GradientCollector;
 
-import java.util.concurrent.atomic.AtomicBoolean;
-
 /** {@code PtGradientCollector} is the PyTorch implementation of {@link GradientCollector}. */
-public final class PtGradientCollector implements GradientCollector {
+public class PtGradientCollector implements GradientCollector {
 
     private boolean gradModel;
-    private static AtomicBoolean isCollecting = new AtomicBoolean();
 
     /** Constructs a new {@code PtGradientCollector} instance. */
     public PtGradientCollector() {
         gradModel = JniUtils.isGradMode();
         JniUtils.setGradMode(true);
-
-        boolean wasCollecting = isCollecting.getAndSet(true);
-        if (wasCollecting) {
-            throw new IllegalStateException(
-                    "A PtGradientCollector is already collecting. Only one can be collecting at a"
-                            + " time");
-        }
-
-        zeroGradients();
     }
 
     /** {@inheritDoc} */
@@ -67,24 +54,12 @@ public final class PtGradientCollector implements GradientCollector {
         JniUtils.backward((PtNDArray) target, (PtNDArray) grad, keepGraph, createGraph);
     }
 
-    /** {@inheritDoc} */
-    @Override
-    public void zeroGradients() {
-        NDManager systemManager = PtNDManager.getSystemManager();
-        for (NDArray array : systemManager.getManagedArrays()) {
-            if (array.hasGradient()) {
-                array.getGradient().subi(array.getGradient());
-            }
-        }
-    }
-
     /** {@inheritDoc} */
     @Override
     public void close() {
         if (!gradModel) {
             JniUtils.setGradMode(false);
         }
-        isCollecting.set(false);
         // TODO: do some clean up if necessary
     }
 }
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 875be4512..9c3a12f7c 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
@@ -175,6 +175,10 @@ public final class JniUtils {
         return new PtNDArray(manager, handle, data);
     }
 
+    public static void emptyCudaCache() {
+        PyTorchLibrary.LIB.torchCudaEmptyCache();
+    }
+
     public static PtNDArray createEmptyNdArray(
             PtNDManager manager, Shape shape, DataType dType, Device device, SparseFormat fmt) {
         int layoutVal = layoutMapper(fmt, device);
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 c28910012..ef48eed33 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
@@ -68,6 +68,8 @@ final class PyTorchLibrary {
 
     native long tensorClone(long handle);
 
+    native void torchCudaEmptyCache();
+
     native long torchEmpty(long[] shape, int dType, int layout, int[] device, boolean requiredGrad);
 
     native long torchZeros(long[] shape, int dType, int layout, int[] device, boolean requiredGrad);
diff --git a/engines/pytorch/pytorch-native/CMakeLists.txt b/engines/pytorch/pytorch-native/CMakeLists.txt
index 8e9c863e8..4453186be 100644
--- a/engines/pytorch/pytorch-native/CMakeLists.txt
+++ b/engines/pytorch/pytorch-native/CMakeLists.txt
@@ -55,6 +55,10 @@ if(PT_VERSION)
     add_compile_definitions(${PT_VERSION})
 endif()
 
+if(USE_CUDA)
+    add_compile_definitions(USE_CUDA)
+endif()
+
 add_library(djl_torch SHARED ${SOURCE_FILES})
 # build host
 if(NOT BUILD_ANDROID)
diff --git a/engines/pytorch/pytorch-native/build.cmd b/engines/pytorch/pytorch-native/build.cmd
index f5ab38ed0..256636c61 100644
--- a/engines/pytorch/pytorch-native/build.cmd
+++ b/engines/pytorch/pytorch-native/build.cmd
@@ -23,6 +23,10 @@ if "%VERSION%" == "1.11.0" (
     set PT_VERSION=V1_11_X
 )
 
+if /i "%2:~0,2%" == "cu" (
+    set USE_CUDA=1
+)
+
 copy /y src\main\patch\cuda.cmake libtorch\share\cmake\Caffe2\public\
 
 @rem workaround VS 17.4.0 issue: https://stackoverflow.com/questions/74366357/updating-to-visual-studio-17-4-0-yields-linker-errors-related-to-tls
@@ -32,5 +36,5 @@ if exist build rd /q /s build
 md build\classes
 cd build
 javac -sourcepath ..\..\pytorch-engine\src\main\java\ ..\..\pytorch-engine\src\main\java\ai\djl\pytorch\jni\PyTorchLibrary.java -h include -d classes
-cmake -DCMAKE_PREFIX_PATH=libtorch -DPT_VERSION=%PT_VERSION% ..
+cmake -DCMAKE_PREFIX_PATH=libtorch -DPT_VERSION=%PT_VERSION% -DUSE_CUDA=%USE_CUDA% ..
 cmake --build . --config Release
diff --git a/engines/pytorch/pytorch-native/build.sh b/engines/pytorch/pytorch-native/build.sh
index 3957008f5..dea17e6f3 100755
--- a/engines/pytorch/pytorch-native/build.sh
+++ b/engines/pytorch/pytorch-native/build.sh
@@ -50,13 +50,17 @@ if [[ "$VERSION" == "1.11.0" ]]; then
   PT_VERSION=V1_11_X
 fi
 
+if [[ "$FLAVOR" = cu* ]]; then
+  USE_CUDA=1
+fi
+
 pushd .
 
 rm -rf build
 mkdir build && cd build
 mkdir classes
 javac -sourcepath ../../pytorch-engine/src/main/java/ ../../pytorch-engine/src/main/java/ai/djl/pytorch/jni/PyTorchLibrary.java -h include -d classes
-cmake -DCMAKE_PREFIX_PATH=libtorch -DPT_VERSION=${PT_VERSION} ..
+cmake -DCMAKE_PREFIX_PATH=libtorch -DPT_VERSION=${PT_VERSION} -DUSE_CUDA=$USE_CUDA ..
 cmake --build . --config Release -- -j "${NUM_PROC}"
 
 if [[ $PLATFORM == 'darwin' ]]; then
diff --git a/engines/pytorch/pytorch-native/src/main/native/ai_djl_pytorch_jni_PyTorchLibrary_system.cc b/engines/pytorch/pytorch-native/src/main/native/ai_djl_pytorch_jni_PyTorchLibrary_system.cc
index e464d579a..505404c0b 100644
--- a/engines/pytorch/pytorch-native/src/main/native/ai_djl_pytorch_jni_PyTorchLibrary_system.cc
+++ b/engines/pytorch/pytorch-native/src/main/native/ai_djl_pytorch_jni_PyTorchLibrary_system.cc
@@ -28,6 +28,10 @@
 #include "djl_pytorch_utils.h"
 #include "ai_djl_pytorch_jni_cache.h"
 
+#ifdef USE_CUDA
+#include <c10/cuda/CUDACachingAllocator.h>
+#endif
+
 #if defined(__ANDROID__)
 #ifndef USE_PTHREADPOOL
 #define USE_PTHREADPOOL
@@ -327,3 +331,13 @@ JNIEXPORT void JNICALL Java_ai_djl_pytorch_jni_PyTorchLibrary_torchStopProfile(
   WriteProfilerEventsToStream(file, event_ptr_lists);
   API_END()
 }
+
+JNIEXPORT void JNICALL Java_ai_djl_pytorch_jni_PyTorchLibrary_torchCudaEmptyCache(JNIEnv* env, jobject jthis) {
+  API_BEGIN()
+#ifdef USE_CUDA
+  if (torch::cuda::is_available()) {
+    c10::cuda::CUDACachingAllocator::emptyCache();
+  }
+#endif
+  API_END()
+}
diff --git a/engines/tensorflow/tensorflow-engine/src/main/java/ai/djl/tensorflow/engine/TfEngine.java b/engines/tensorflow/tensorflow-engine/src/main/java/ai/djl/tensorflow/engine/TfEngine.java
index bd8d10f36..631194bc0 100644
--- a/engines/tensorflow/tensorflow-engine/src/main/java/ai/djl/tensorflow/engine/TfEngine.java
+++ b/engines/tensorflow/tensorflow-engine/src/main/java/ai/djl/tensorflow/engine/TfEngine.java
@@ -21,8 +21,6 @@ import ai.djl.ndarray.NDManager;
 import ai.djl.nn.SymbolBlock;
 import ai.djl.tensorflow.engine.javacpp.JavacppUtils;
 import ai.djl.tensorflow.engine.javacpp.LibUtils;
-import ai.djl.training.GradientCollector;
-import ai.djl.util.RandomUtils;
 
 import org.bytedeco.javacpp.PointerScope;
 import org.tensorflow.TensorFlow;
@@ -161,19 +159,6 @@ public final class TfEngine extends Engine implements AutoCloseable {
         return TfNDManager.getSystemManager().newSubManager(device);
     }
 
-    /** {@inheritDoc} */
-    @Override
-    public GradientCollector newGradientCollector() {
-        throw new UnsupportedOperationException("TensorFlow does not support training yet");
-    }
-
-    /** {@inheritDoc} */
-    @Override
-    public void setRandomSeed(int seed) {
-        super.setRandomSeed(seed);
-        RandomUtils.RANDOM.setSeed(seed);
-    }
-
     TFE_Context getEagerSession() {
         return eagerSessionHandle.get();
     }
diff --git a/engines/tensorrt/src/main/java/ai/djl/tensorrt/engine/TrtEngine.java b/engines/tensorrt/src/main/java/ai/djl/tensorrt/engine/TrtEngine.java
index e56633552..3e58f4ab9 100644
--- a/engines/tensorrt/src/main/java/ai/djl/tensorrt/engine/TrtEngine.java
+++ b/engines/tensorrt/src/main/java/ai/djl/tensorrt/engine/TrtEngine.java
@@ -17,11 +17,8 @@ import ai.djl.Model;
 import ai.djl.engine.Engine;
 import ai.djl.engine.EngineException;
 import ai.djl.engine.StandardCapabilities;
-import ai.djl.ndarray.NDManager;
-import ai.djl.nn.SymbolBlock;
 import ai.djl.tensorrt.jni.JniUtils;
 import ai.djl.tensorrt.jni.LibUtils;
-import ai.djl.training.GradientCollector;
 import ai.djl.util.Utils;
 
 import java.io.FileNotFoundException;
@@ -106,12 +103,6 @@ public final class TrtEngine extends Engine {
         return StandardCapabilities.CUDA.equals(capability);
     }
 
-    /** {@inheritDoc} */
-    @Override
-    public SymbolBlock newSymbolBlock(NDManager manager) {
-        throw new UnsupportedOperationException("TensorRT does not support empty SymbolBlock");
-    }
-
     /** {@inheritDoc} */
     @Override
     public Model newModel(String name, Device device) {
@@ -135,18 +126,6 @@ public final class TrtEngine extends Engine {
         return TrtNDManager.getSystemManager().newSubManager(device);
     }
 
-    /** {@inheritDoc} */
-    @Override
-    public GradientCollector newGradientCollector() {
-        throw new UnsupportedOperationException("Not supported for TensorRT");
-    }
-
-    /** {@inheritDoc} */
-    @Override
-    public void setRandomSeed(int seed) {
-        throw new UnsupportedOperationException("Not supported for TensorRT");
-    }
-
     /** {@inheritDoc} */
     @Override
     public String toString() {
diff --git a/engines/tflite/tflite-engine/src/main/java/ai/djl/tflite/engine/TfLiteEngine.java b/engines/tflite/tflite-engine/src/main/java/ai/djl/tflite/engine/TfLiteEngine.java
index 02be61faa..731db9a1e 100644
--- a/engines/tflite/tflite-engine/src/main/java/ai/djl/tflite/engine/TfLiteEngine.java
+++ b/engines/tflite/tflite-engine/src/main/java/ai/djl/tflite/engine/TfLiteEngine.java
@@ -17,8 +17,6 @@ import ai.djl.Device;
 import ai.djl.Model;
 import ai.djl.engine.Engine;
 import ai.djl.ndarray.NDManager;
-import ai.djl.nn.SymbolBlock;
-import ai.djl.training.GradientCollector;
 
 /**
  * The {@code TfLiteEngine} is an implementation of the {@link Engine} based on the <a
@@ -89,12 +87,6 @@ public final class TfLiteEngine extends Engine {
         return new TfLiteModel(name, newBaseManager(device));
     }
 
-    /** {@inheritDoc} */
-    @Override
-    public SymbolBlock newSymbolBlock(NDManager manager) {
-        throw new UnsupportedOperationException("TFLite does not support empty SymbolBlock");
-    }
-
     /** {@inheritDoc} */
     @Override
     public NDManager newBaseManager() {
@@ -106,16 +98,4 @@ public final class TfLiteEngine extends Engine {
     public NDManager newBaseManager(Device device) {
         return TfLiteNDManager.getSystemManager().newSubManager(device);
     }
-
-    /** {@inheritDoc} */
-    @Override
-    public GradientCollector newGradientCollector() {
-        throw new UnsupportedOperationException("Not supported for TFLite");
-    }
-
-    /** {@inheritDoc} */
-    @Override
-    public void setRandomSeed(int seed) {
-        throw new UnsupportedOperationException("Not supported for TFLite");
-    }
 }


Test patch:
diff --git a/engines/pytorch/pytorch-engine/src/test/java/ai/djl/pytorch/jni/JniUtilsTest.java b/engines/pytorch/pytorch-engine/src/test/java/ai/djl/pytorch/jni/JniUtilsTest.java
new file mode 100644
index 000000000..0e7a6ae98
--- /dev/null
+++ b/engines/pytorch/pytorch-engine/src/test/java/ai/djl/pytorch/jni/JniUtilsTest.java
@@ -0,0 +1,23 @@
+/*
+ * Copyright 2023 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.pytorch.jni;
+
+import org.testng.annotations.Test;
+
+public class JniUtilsTest {
+
+    @Test
+    public void testClearGpuCache() {
+        JniUtils.emptyCudaCache();
+    }
+}
diff --git a/examples/src/test/java/ai/djl/examples/training/TrainAirfoilWithTabNetTest.java b/examples/src/test/java/ai/djl/examples/training/TrainAirfoilWithTabNetTest.java
index 403ca2c50..89a2dadcc 100644
--- a/examples/src/test/java/ai/djl/examples/training/TrainAirfoilWithTabNetTest.java
+++ b/examples/src/test/java/ai/djl/examples/training/TrainAirfoilWithTabNetTest.java
@@ -24,7 +24,6 @@ import java.io.IOException;
 public class TrainAirfoilWithTabNetTest {
     @Test
     public void testTrainAirfoilWithTabNet() throws TranslateException, IOException {
-        TestRequirements.nightly();
         TestRequirements.engine("MXNet", "PyTorch");
         String[] args = new String[] {"-g", "1", "-e", "20", "-b", "32"};
         TrainingResult result = TrainAirfoilWithTabNet.runExample(args);
diff --git a/integration/src/main/java/ai/djl/integration/tests/training/GradientCollectorIntegrationTest.java b/integration/src/main/java/ai/djl/integration/tests/training/GradientCollectorIntegrationTest.java
index 2ad4f5e3b..d8469178d 100644
--- a/integration/src/main/java/ai/djl/integration/tests/training/GradientCollectorIntegrationTest.java
+++ b/integration/src/main/java/ai/djl/integration/tests/training/GradientCollectorIntegrationTest.java
@@ -82,90 +82,25 @@ public class GradientCollectorIntegrationTest {
         }
     }
 
-    @Test
-    public void testZeroGradients() {
-        try (NDManager manager = NDManager.newBaseManager(TestUtils.getEngine())) {
-            NDArray a = manager.create(0.0f);
-            a.setRequiresGradient(true);
-
-            Engine engine = Engine.getEngine(TestUtils.getEngine());
-            try (GradientCollector gc = engine.newGradientCollector()) {
-                NDArray b = a.mul(2);
-
-                // Gradients are initially zero
-                Assert.assertEquals(a.getGradient().getFloat(), 0.0f);
-
-                // Gradients are updated by backwards
-                gc.backward(b);
-                Assert.assertEquals(a.getGradient().getFloat(), 2.0f);
-
-                // Gradients are cleared by zeroGradients
-                gc.zeroGradients();
-                Assert.assertEquals(a.getGradient().getFloat(), 0.0f);
-            }
-        }
-    }
-
     /** Tests that the gradients do not accumulate when closing the gradient collector. */
     @Test
     public void testClearGradients() {
         try (NDManager manager = NDManager.newBaseManager(TestUtils.getEngine())) {
-            NDArray a = manager.create(0.0f);
-            a.setRequiresGradient(true);
+            NDArray variable = manager.create(0.0f);
+            variable.setRequiresGradient(true);
 
-            Engine engine = Engine.getEngine(TestUtils.getEngine());
+            Engine engine = manager.getEngine();
             for (int i = 0; i < 3; i++) {
+                manager.zeroGradients();
                 try (GradientCollector gc = engine.newGradientCollector()) {
-                    NDArray b = a.mul(2);
-                    gc.backward(b);
-                }
-                Assert.assertEquals(a.getGradient().getFloat(), 2.0f);
-            }
-        }
-    }
-
-    /** Tests that the gradients do accumulate within the same gradient collector. */
-    @Test
-    public void testAccumulateGradients() {
-        // TODO: MXNet support for accumulating gradients does not currently work
-        TestRequirements.notEngine("MXNet");
-        try (NDManager manager = NDManager.newBaseManager(TestUtils.getEngine())) {
-            NDArray a = manager.create(0.0f);
-            a.setRequiresGradient(true);
-
-            Engine engine = Engine.getEngine(TestUtils.getEngine());
-            try (GradientCollector gc = engine.newGradientCollector()) {
-                for (int i = 1; i <= 3; i++) {
-                    NDArray b = a.mul(2);
-                    gc.backward(b);
-                    Assert.assertEquals(a.getGradient().getFloat(), 2.0f * i);
+                    NDArray loss = variable.mul(2);
+                    gc.backward(loss);
                 }
+                Assert.assertEquals(variable.getGradient().getFloat(), 2.0f);
             }
         }
     }
 
-    /**
-     * Ensures that a gradient collector does not start when one is already created because they are
-     * global.
-     */
-    @Test
-    @SuppressWarnings({"try", "PMD.UseTryWithResources"})
-    public void testMultipleGradientCollectors() {
-        Assert.assertThrows(
-                () -> {
-                    GradientCollector gc2 = null;
-                    Engine engine = Engine.getEngine(TestUtils.getEngine());
-                    try (GradientCollector gc = engine.newGradientCollector()) {
-                        gc2 = engine.newGradientCollector();
-                        gc2.close();
-                    } finally {
-                        if (gc2 != null) {
-                            gc2.close();
-                        }
-                    }
-                });
-    }
-
     @Test
     public void testFreezeParameters() {
         try (Model model = Model.newInstance("model", TestUtils.getEngine())) {