Repository: deepjavalibrary/djl
Pull request #2304 (2023-01-07 04:03:52)

A temporary solution to issue 2210 
This PR aims to solve the issue https://github.com/deepjavalibrary/djl/issues/2210.  As mentioned in https://github.com/deepjavalibrary/djl/pull/2232, the issue originates from `zeroGradients()` call, which was added in PR https://github.com/deepjavalibrary/djl/pull/2101 in order to solve issue https://github.com/deepjavalibrary/djl/issues/2024. As attempted in PR https://github.com/deepjavalibrary/djl/pull/2232, specifying the block to a PtGradientCollector does not really solve the problem, but it is more of adding a feature. In that PR, by default, creation of PtGradientCollector doesn't call zeroGradients(), which means the zeroGradients feature introduced in PR https://github.com/deepjavalibrary/djl/pull/2101 has been effectively reverted in the solution proposed there.

Since there is currently no well designed solution to the efficiency issue mentioned in https://github.com/deepjavalibrary/djl/issues/2210, here the temporary solution is to revert PR 2111 and PR2101. To solve the original issue https://github.com/deepjavalibrary/djl/issues/2024 in the d2l book, we can just modify the example in the book to manually zero gradient, as shown below. The idea is that currently the GradientCollector does not do any additional modification after collecting the gradient from the engines. So it keeps default behaviour of each engine in terms of the gradient aggregation. Ie, for mxnet gradients are replaced, while for pytorch engine the gradients are added. This approach also maintains the room to add features of specifying how gradient is accumulated in the future. See https://github.com/deepjavalibrary/djl/pull/2232.

---
```java
    public void testClearGradients() {
        try (NDManager manager = NDManager.newBaseManager()) {
            NDArray variable = manager.create(0.0f);
            variable.setRequiresGradient(true);

            for (int i = 0; i < 3; i++) {
                manager.zeroGradients();
                try (GradientCollector gc = Engine.getInstance().newGradientCollector()) {
                    NDArray loss = variable.mul(2);
                    gc.backward(loss);
                }
                Assert.assertEquals(variable.getGradient().getFloat(), 2.0f);
            }
        }
    }
```


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

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


```diff
@@             Coverage Diff              @@
##             master    #2304      +/-   ##
============================================
+ Coverage     72.08%   74.20%   +2.11%     
- Complexity     5126     6674    +1548     
============================================
  Files           473      656     +183     
  Lines         21970    29054    +7084     
  Branches       2351     2999     +648     
============================================
+ Hits          15838    21560    +5722     
- Misses         4925     6031    +1106     
- Partials       1207     1463     +256     
```


| [Impacted Files](https://codecov.io/gh/deepjavalibrary/djl/pull/2304?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/2304?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/2304?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/2304?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/2304?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/2304?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/2304?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/2304?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/2304?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/2304?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/2304?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 [598 more](https://codecov.io/gh/deepjavalibrary/djl/pull/2304?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/2304?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/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/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/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
     }
 }


Test patch:
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())) {