Repository: deepjavalibrary/djl
Pull request #2293 (2023-01-05 04:59:52)

add openai whisper model
## Description ##

Add OpenAI Whisper model to DJL


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/modality/audio/Audio.java b/api/src/main/java/ai/djl/modality/audio/Audio.java
index 4ad949845..429359fc8 100644
--- a/api/src/main/java/ai/djl/modality/audio/Audio.java
+++ b/api/src/main/java/ai/djl/modality/audio/Audio.java
@@ -63,9 +63,9 @@ public class Audio {
     }
 
     /**
-     * Get the number of channels of an audio file.
+     * Returns the number of channels of an audio file.
      *
-     * @return The number of channels of an audio file.
+     * @return the number of channels of an audio file
      */
     public int getChannels() {
         return channels;
diff --git a/api/src/main/java/ai/djl/modality/audio/AudioFactory.java b/api/src/main/java/ai/djl/modality/audio/AudioFactory.java
index 310134c57..97d6446c7 100644
--- a/api/src/main/java/ai/djl/modality/audio/AudioFactory.java
+++ b/api/src/main/java/ai/djl/modality/audio/AudioFactory.java
@@ -36,9 +36,16 @@ public abstract class AudioFactory {
         "ai.djl.audio.FFmpegAudioFactory", "ai.djl.modality.audio.SampledAudioFactory"
     };
 
-    private static AudioFactory factory = newInstance();
+    protected int channels;
+    protected int sampleRate;
+    protected int sampleFormat;
 
-    private static AudioFactory newInstance() {
+    /**
+     * Constructs a new instance of {@code AudioFactory}.
+     *
+     * @return a new instance of {@code AudioFactory}
+     */
+    public static AudioFactory newInstance() {
         for (String f : FACTORIES) {
             try {
                 Class<? extends AudioFactory> clazz =
@@ -51,15 +58,6 @@ public abstract class AudioFactory {
         throw new IllegalStateException("Failed to create AudioFactory!");
     }
 
-    /**
-     * Gets new instance of Audio factory from the provided factory implementation.
-     *
-     * @return {@link AudioFactory}
-     */
-    public static AudioFactory getInstance() {
-        return factory;
-    }
-
     /**
      * Returns {@link Audio} from file.
      *
@@ -112,7 +110,9 @@ public abstract class AudioFactory {
      * @param data the raw data in float array form.
      * @return {@link Audio}
      */
-    public abstract Audio fromData(float[] data);
+    public Audio fromData(float[] data) {
+        return new Audio(data);
+    }
 
     /**
      * Returns {@link Audio} from {@link NDArray}.
@@ -120,5 +120,67 @@ public abstract class AudioFactory {
      * @param array the NDArray with CHW format
      * @return {@link Audio}
      */
-    public abstract Audio fromNDArray(NDArray array);
+    public Audio fromNDArray(NDArray array) {
+        throw new UnsupportedOperationException("Not supported!");
+    }
+
+    /**
+     * Sets the number of channels for {@link AudioFactory} to use.
+     *
+     * @param channels the number of channels for {@link AudioFactory} to use
+     * @return this factory
+     */
+    public AudioFactory setChannels(int channels) {
+        this.channels = channels;
+        return this;
+    }
+
+    /**
+     * Returns the channels of this factory.
+     *
+     * @return the channels of this factory
+     */
+    public int getChannels() {
+        return channels;
+    }
+
+    /**
+     * Sets the sampleRate for {@link AudioFactory} to use.
+     *
+     * @param sampleRate the sampleRate for {@link AudioFactory} to use
+     * @return this factory
+     */
+    public AudioFactory setSampleRate(int sampleRate) {
+        this.sampleRate = sampleRate;
+        return this;
+    }
+
+    /**
+     * Returns the sample rate.
+     *
+     * @return the sample rate in integer
+     */
+    public int getSampleRate() {
+        return sampleRate;
+    }
+
+    /**
+     * Sets the audio sample format for {@link AudioFactory} to use.
+     *
+     * @param sampleFormat the sample format
+     * @return this factory.
+     */
+    public AudioFactory setSampleFormat(int sampleFormat) {
+        this.sampleFormat = sampleFormat;
+        return this;
+    }
+
+    /**
+     * Returns the sample format name of the audio.
+     *
+     * @return the format name of the audio
+     */
+    public int getSampleFormat() {
+        return sampleFormat;
+    }
 }
diff --git a/api/src/main/java/ai/djl/modality/audio/SampledAudioFactory.java b/api/src/main/java/ai/djl/modality/audio/SampledAudioFactory.java
index cdbeaccf2..04ce1691f 100644
--- a/api/src/main/java/ai/djl/modality/audio/SampledAudioFactory.java
+++ b/api/src/main/java/ai/djl/modality/audio/SampledAudioFactory.java
@@ -13,7 +13,6 @@
 package ai.djl.modality.audio;
 
 import ai.djl.modality.cv.ImageFactory;
-import ai.djl.ndarray.NDArray;
 
 import java.io.BufferedInputStream;
 import java.io.IOException;
@@ -37,16 +36,34 @@ import javax.sound.sampled.UnsupportedAudioFileException;
  */
 public class SampledAudioFactory extends AudioFactory {
 
+    /** {@inheritDoc} */
+    @Override
+    public AudioFactory setChannels(int channel) {
+        if (channel != 0) {
+            throw new UnsupportedOperationException("SampledAudioFactory only support 1 channel.");
+        }
+        return this;
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public AudioFactory setSampleRate(int sampleRate) {
+        throw new UnsupportedOperationException("Not supported.");
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public AudioFactory setSampleFormat(int sampleFormat) {
+        throw new UnsupportedOperationException("Not supported.");
+    }
+
     /** {@inheritDoc} */
     @Override
     public Audio fromFile(Path path) throws IOException {
         try (AudioInputStream ais = AudioSystem.getAudioInputStream(path.toFile())) {
             AudioFormat format = ais.getFormat();
             byte[] bytes = read(ais);
-            float[] floats =
-                    bytesToFloats(
-                            bytes,
-                            format.isBigEndian() ? ByteOrder.BIG_ENDIAN : ByteOrder.LITTLE_ENDIAN);
+            float[] floats = bytesToFloats(bytes, format.isBigEndian());
             return new Audio(floats, format.getSampleRate(), format.getChannels());
         } catch (UnsupportedAudioFileException e) {
             throw new IOException("Unsupported Audio file", e);
@@ -59,28 +76,13 @@ public class SampledAudioFactory extends AudioFactory {
         try (AudioInputStream ais = AudioSystem.getAudioInputStream(new BufferedInputStream(is))) {
             AudioFormat format = ais.getFormat();
             byte[] bytes = read(ais);
-            float[] floats =
-                    bytesToFloats(
-                            bytes,
-                            format.isBigEndian() ? ByteOrder.BIG_ENDIAN : ByteOrder.LITTLE_ENDIAN);
+            float[] floats = bytesToFloats(bytes, format.isBigEndian());
             return new Audio(floats, format.getSampleRate(), format.getChannels());
         } catch (UnsupportedAudioFileException e) {
             throw new IOException("Unsupported Audio file", e);
         }
     }
 
-    /** {@inheritDoc} */
-    @Override
-    public Audio fromData(float[] data) {
-        return new Audio(data);
-    }
-
-    /** {@inheritDoc} */
-    @Override
-    public Audio fromNDArray(NDArray array) {
-        throw new UnsupportedOperationException("Not supported!");
-    }
-
     private byte[] read(AudioInputStream ais) throws IOException {
         AudioFormat format = ais.getFormat();
         int frameSize = format.getFrameSize();
@@ -102,7 +104,8 @@ public class SampledAudioFactory extends AudioFactory {
         return ret;
     }
 
-    private float[] bytesToFloats(byte[] bytes, ByteOrder order) {
+    private float[] bytesToFloats(byte[] bytes, boolean isBigEndian) {
+        ByteOrder order = isBigEndian ? ByteOrder.BIG_ENDIAN : ByteOrder.LITTLE_ENDIAN;
         ShortBuffer buffer = ByteBuffer.wrap(bytes).order(order).asShortBuffer();
         short[] shorts = new short[buffer.capacity()];
         buffer.get(shorts);
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/nn/norm/LayerNorm.java b/api/src/main/java/ai/djl/nn/norm/LayerNorm.java
index 43ebd85b8..5d6928413 100644
--- a/api/src/main/java/ai/djl/nn/norm/LayerNorm.java
+++ b/api/src/main/java/ai/djl/nn/norm/LayerNorm.java
@@ -178,7 +178,7 @@ public class LayerNorm extends AbstractBlock {
     }
 
     /** The Builder to construct a {@link LayerNorm}. */
-    public static final class Builder {
+    public static class Builder {
 
         private float epsilon = 1E-5f;
         // private Shape normalizedShape;
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/docker/spark/Dockerfile b/docker/spark/Dockerfile
new file mode 100644
index 000000000..a0b346633
--- /dev/null
+++ b/docker/spark/Dockerfile
@@ -0,0 +1,31 @@
+# -*- mode: dockerfile -*-
+# 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.txt" file accompanying this file. This file is distributed on an "AS IS"
+# BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, express or implied. See the License for
+# the specific language governing permissions and limitations under the License.
+FROM 895885662937.dkr.ecr.us-west-2.amazonaws.com/spark/emr-6.9.0:latest
+LABEL maintainer="djl-dev@amazon.com"
+
+# Add DJL jars
+USER root
+ARG DJL_VERSION=0.20.0
+ARG JNA_VERSION=5.12.1
+ADD https://repo1.maven.org/maven2/ai/djl/api/${DJL_VERSION}/api-${DJL_VERSION}.jar /usr/lib/spark/jars/djl-api-${DJL_VERSION}.jar
+ADD https://repo1.maven.org/maven2/ai/djl/spark/spark/${DJL_VERSION}/spark-${DJL_VERSION}.jar /usr/lib/spark/jars/djl-spark-${DJL_VERSION}.jar
+ADD https://repo1.maven.org/maven2/ai/djl/pytorch/pytorch-engine/${DJL_VERSION}/pytorch-engine-${DJL_VERSION}.jar /usr/lib/spark/jars/djl-pytorch-engine-${DJL_VERSION}.jar
+ADD https://repo1.maven.org/maven2/ai/djl/mxnet/mxnet-engine/${DJL_VERSION}/mxnet-engine-${DJL_VERSION}.jar /usr/lib/spark/jars/djl-mxnet-engine-${DJL_VERSION}.jar
+ADD https://repo1.maven.org/maven2/ai/djl/onnxruntime/onnxruntime-engine/${DJL_VERSION}/onnxruntime-engine-${DJL_VERSION}.jar /usr/lib/spark/jars/djl-onnxruntime-engine-${DJL_VERSION}.jar
+ADD https://repo1.maven.org/maven2/ai/djl/tensorflow/tensorflow-engine/${DJL_VERSION}/tensorflow-engine-${DJL_VERSION}.jar /usr/lib/spark/jars/djl-tensorflow-engine-${DJL_VERSION}.jar
+ADD https://repo1.maven.org/maven2/net/java/dev/jna/jna/${JNA_VERSION}/jna-${JNA_VERSION}.jar /usr/lib/spark/jars/
+RUN chmod -R +r /usr/lib/spark/jars/
+
+# Set environment
+ENV PYTORCH_PRECXX11 true
+
+USER hadoop:hadoop
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/mxnet/mxnet-engine/src/main/java/ai/djl/mxnet/jna/LibUtils.java b/engines/mxnet/mxnet-engine/src/main/java/ai/djl/mxnet/jna/LibUtils.java
index 8dc9ca6b3..793f14691 100644
--- a/engines/mxnet/mxnet-engine/src/main/java/ai/djl/mxnet/jna/LibUtils.java
+++ b/engines/mxnet/mxnet-engine/src/main/java/ai/djl/mxnet/jna/LibUtils.java
@@ -296,6 +296,7 @@ public final class LibUtils {
     private static boolean supported(Platform platform) {
         // mxnet-native-cu102mkl:1.8.0: 3.0, 5.0, 6.0, 7.0, 7.5
         // mxnet-native-cu110mkl:1.8.0: 5.0, 6.0, 7.0, 8.0
+        // mxnet-native-cu112mkl:1.9.1: 5.0, 6.0, 7.0, 7.5, 8.0, 8.6
         String version = platform.getVersion();
         if (version.startsWith("1.8.") || version.startsWith("1.9.")) {
             String flavor = platform.getFlavor();
@@ -304,7 +305,7 @@ public final class LibUtils {
                 if (version.startsWith("1.8.")) {
                     return Arrays.asList("50", "60", "70", "80").contains(cudaArch);
                 }
-                return Arrays.asList("50", "60", "70", "75", "80").contains(cudaArch);
+                return Arrays.asList("50", "60", "70", "75", "80", "86").contains(cudaArch);
             } else if (flavor.startsWith("cu10")) {
                 return Arrays.asList("30", "50", "60", "70", "75").contains(cudaArch);
             }
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 f0b2d1831..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);
@@ -364,57 +368,63 @@ public final class JniUtils {
         }
         List<NDIndexElement> indices = index.getIndices();
         long torchIndexHandle = PyTorchLibrary.LIB.torchIndexInit(indices.size());
-        ListIterator<NDIndexElement> it = indices.listIterator();
-        while (it.hasNext()) {
-            if (it.nextIndex() == index.getEllipsisIndex()) {
-                PyTorchLibrary.LIB.torchIndexAppendNoneEllipsis(torchIndexHandle, true);
-            }
+        try {
+            // Index aggregation
+            ListIterator<NDIndexElement> it = indices.listIterator();
+            while (it.hasNext()) {
+                if (it.nextIndex() == index.getEllipsisIndex()) {
+                    PyTorchLibrary.LIB.torchIndexAppendNoneEllipsis(torchIndexHandle, true);
+                }
 
-            NDIndexElement elem = it.next();
-            if (elem instanceof NDIndexNull) {
-                PyTorchLibrary.LIB.torchIndexAppendNoneEllipsis(torchIndexHandle, false);
-            } else if (elem instanceof NDIndexSlice) {
-                Long min = ((NDIndexSlice) elem).getMin();
-                Long max = ((NDIndexSlice) elem).getMax();
-                Long step = ((NDIndexSlice) elem).getStep();
-                int nullSliceBin = (min == null ? 1 : 0) * 2 + (max == null ? 1 : 0);
-                // nullSliceBin encodes whether the slice (min, max) is null:
-                // is_null == 1, ! is_null == 0;
-                // 0b11 == 3, 0b10 = 2, ...
-                PyTorchLibrary.LIB.torchIndexAppendSlice(
-                        torchIndexHandle,
-                        min == null ? 0 : min,
-                        max == null ? 0 : max,
-                        step == null ? 1 : step,
-                        nullSliceBin);
-            } else if (elem instanceof NDIndexAll) {
-                PyTorchLibrary.LIB.torchIndexAppendSlice(torchIndexHandle, 0, 0, 1, 3);
-            } else if (elem instanceof NDIndexFixed) {
-                PyTorchLibrary.LIB.torchIndexAppendFixed(
-                        torchIndexHandle, ((NDIndexFixed) elem).getIndex());
-            } else if (elem instanceof NDIndexBooleans) {
-                PtNDArray indexArr = (PtNDArray) ((NDIndexBooleans) elem).getIndex();
-                PyTorchLibrary.LIB.torchIndexAppendArray(torchIndexHandle, indexArr.getHandle());
-            } else if (elem instanceof NDIndexTake) {
-                PtNDArray indexArr = manager.from(((NDIndexTake) elem).getIndex());
-                if (indexArr.getDataType() != DataType.INT64) {
-                    indexArr = indexArr.toType(DataType.INT64, true);
+                NDIndexElement elem = it.next();
+                if (elem instanceof NDIndexNull) {
+                    PyTorchLibrary.LIB.torchIndexAppendNoneEllipsis(torchIndexHandle, false);
+                } else if (elem instanceof NDIndexSlice) {
+                    Long min = ((NDIndexSlice) elem).getMin();
+                    Long max = ((NDIndexSlice) elem).getMax();
+                    Long step = ((NDIndexSlice) elem).getStep();
+                    int nullSliceBinary = (min == null ? 1 : 0) * 2 + (max == null ? 1 : 0);
+                    // nullSliceBinary encodes whether the slice end {min, max} is null:
+                    // is_null == 1, ! is_null == 0;
+                    // 0b11 == 3, 0b10 = 2, ...
+                    // If {min, max} is null, then its value is ineffective, thus set to -1.
+                    PyTorchLibrary.LIB.torchIndexAppendSlice(
+                            torchIndexHandle,
+                            min == null ? -1 : min,
+                            max == null ? -1 : max,
+                            step == null ? 1 : step,
+                            nullSliceBinary);
+                } else if (elem instanceof NDIndexAll) {
+                    PyTorchLibrary.LIB.torchIndexAppendSlice(torchIndexHandle, -1, -1, 1, 3);
+                } else if (elem instanceof NDIndexFixed) {
+                    PyTorchLibrary.LIB.torchIndexAppendFixed(
+                            torchIndexHandle, ((NDIndexFixed) elem).getIndex());
+                } else if (elem instanceof NDIndexBooleans) {
+                    PtNDArray indexArr = (PtNDArray) ((NDIndexBooleans) elem).getIndex();
+                    PyTorchLibrary.LIB.torchIndexAppendArray(
+                            torchIndexHandle, indexArr.getHandle());
+                } else if (elem instanceof NDIndexTake) {
+                    PtNDArray indexArr = manager.from(((NDIndexTake) elem).getIndex());
+                    if (indexArr.getDataType() != DataType.INT64) {
+                        indexArr = indexArr.toType(DataType.INT64, true);
+                    }
+                    PyTorchLibrary.LIB.torchIndexAppendArray(
+                            torchIndexHandle, indexArr.getHandle());
+                } else if (elem instanceof NDIndexPick) {
+                    // Backward compatible
+                    NDIndexFullPick fullPick =
+                            NDIndexFullPick.fromIndex(index, ndArray.getShape()).get();
+                    return pick(ndArray, manager.from(fullPick.getIndices()), fullPick.getAxis());
                 }
-                PyTorchLibrary.LIB.torchIndexAppendArray(torchIndexHandle, indexArr.getHandle());
-            } else if (elem instanceof NDIndexPick) {
-                // Backward compatible
-                NDIndexFullPick fullPick =
-                        NDIndexFullPick.fromIndex(index, ndArray.getShape()).get();
-                return pick(ndArray, manager.from(fullPick.getIndices()), fullPick.getAxis());
             }
+            if (indices.size() == index.getEllipsisIndex()) {
+                PyTorchLibrary.LIB.torchIndexAppendNoneEllipsis(torchIndexHandle, true);
+            }
+            long ret = PyTorchLibrary.LIB.torchIndexAdvGet(ndArray.getHandle(), torchIndexHandle);
+            return new PtNDArray(manager, ret);
+        } finally {
+            PyTorchLibrary.LIB.torchDeleteIndex(torchIndexHandle);
         }
-        if (indices.size() == index.getEllipsisIndex()) {
-            PyTorchLibrary.LIB.torchIndexAppendNoneEllipsis(torchIndexHandle, true);
-        }
-
-        return new PtNDArray(
-                manager,
-                PyTorchLibrary.LIB.torchIndexAdvGet(ndArray.getHandle(), torchIndexHandle));
     }
 
     @SuppressWarnings("OptionalGetWithoutIsPresent")
@@ -422,61 +432,70 @@ public final class JniUtils {
         if (ndArray == null) {
             return;
         }
-
-        // Index aggregation
         List<NDIndexElement> indices = index.getIndices();
         long torchIndexHandle = PyTorchLibrary.LIB.torchIndexInit(indices.size());
-        ListIterator<NDIndexElement> it = indices.listIterator();
-        while (it.hasNext()) {
-            if (it.nextIndex() == index.getEllipsisIndex()) {
-                PyTorchLibrary.LIB.torchIndexAppendNoneEllipsis(torchIndexHandle, true);
-            }
+        try {
+            // Index aggregation
+            ListIterator<NDIndexElement> it = indices.listIterator();
+            while (it.hasNext()) {
+                if (it.nextIndex() == index.getEllipsisIndex()) {
+                    PyTorchLibrary.LIB.torchIndexAppendNoneEllipsis(torchIndexHandle, true);
+                }
 
-            NDIndexElement elem = it.next();
-            if (elem instanceof NDIndexNull) {
-                PyTorchLibrary.LIB.torchIndexAppendNoneEllipsis(torchIndexHandle, false);
-            } else if (elem instanceof NDIndexSlice) {
-                Long min = ((NDIndexSlice) elem).getMin();
-                Long max = ((NDIndexSlice) elem).getMax();
-                Long step = ((NDIndexSlice) elem).getStep();
-                int nullSliceBin = (min == null ? 1 : 0) * 2 + (max == null ? 1 : 0);
-                // nullSliceBin encodes whether the slice (min, max) is null:
-                // is_null == 1, ! is_null == 0;
-                // 0b11 == 3, 0b10 = 2, ...
-                PyTorchLibrary.LIB.torchIndexAppendSlice(
-                        torchIndexHandle,
-                        min == null ? 0 : min,
-                        max == null ? 0 : max,
-                        step == null ? 1 : step,
-                        nullSliceBin);
-            } else if (elem instanceof NDIndexAll) {
-                PyTorchLibrary.LIB.torchIndexAppendSlice(torchIndexHandle, 0, 0, 1, 3);
-            } else if (elem instanceof NDIndexFixed) {
-                PyTorchLibrary.LIB.torchIndexAppendFixed(
-                        torchIndexHandle, ((NDIndexFixed) elem).getIndex());
-            } else if (elem instanceof NDIndexBooleans) {
-                PtNDArray indexArr = (PtNDArray) ((NDIndexBooleans) elem).getIndex();
-                PyTorchLibrary.LIB.torchIndexAppendArray(torchIndexHandle, indexArr.getHandle());
-            } else if (elem instanceof NDIndexTake) {
-                PtNDArray indexArr = (PtNDArray) ((NDIndexTake) elem).getIndex();
-                if (indexArr.getDataType() != DataType.INT64) {
-                    indexArr = indexArr.toType(DataType.INT64, true);
+                NDIndexElement elem = it.next();
+                if (elem instanceof NDIndexNull) {
+                    PyTorchLibrary.LIB.torchIndexAppendNoneEllipsis(torchIndexHandle, false);
+                } else if (elem instanceof NDIndexSlice) {
+                    Long min = ((NDIndexSlice) elem).getMin();
+                    Long max = ((NDIndexSlice) elem).getMax();
+                    Long step = ((NDIndexSlice) elem).getStep();
+                    int nullSliceBinary = (min == null ? 1 : 0) * 2 + (max == null ? 1 : 0);
+                    // nullSliceBinary encodes whether the slice end {min, max} is null:
+                    // is_null == 1, ! is_null == 0;
+                    // 0b11 == 3, 0b10 = 2, ...
+                    // If {min, max} is null, then its value is ineffective, thus set to -1.
+                    PyTorchLibrary.LIB.torchIndexAppendSlice(
+                            torchIndexHandle,
+                            min == null ? -1 : min,
+                            max == null ? -1 : max,
+                            step == null ? 1 : step,
+                            nullSliceBinary);
+                } else if (elem instanceof NDIndexAll) {
+                    PyTorchLibrary.LIB.torchIndexAppendSlice(torchIndexHandle, -1, -1, 1, 3);
+                } else if (elem instanceof NDIndexFixed) {
+                    PyTorchLibrary.LIB.torchIndexAppendFixed(
+                            torchIndexHandle, ((NDIndexFixed) elem).getIndex());
+                } else if (elem instanceof NDIndexBooleans) {
+                    PtNDArray indexArr = (PtNDArray) ((NDIndexBooleans) elem).getIndex();
+                    PyTorchLibrary.LIB.torchIndexAppendArray(
+                            torchIndexHandle, indexArr.getHandle());
+                } else if (elem instanceof NDIndexTake) {
+                    PtNDArray indexArr = (PtNDArray) ((NDIndexTake) elem).getIndex();
+                    if (indexArr.getDataType() != DataType.INT64) {
+                        indexArr = indexArr.toType(DataType.INT64, true);
+                    }
+                    PyTorchLibrary.LIB.torchIndexAppendArray(
+                            torchIndexHandle, indexArr.getHandle());
+                } else if (elem instanceof NDIndexPick) {
+                    // Backward compatible
+                    NDIndexFullPick fullPick =
+                            NDIndexFullPick.fromIndex(index, ndArray.getShape()).get();
+                    pick(
+                            ndArray,
+                            ndArray.getManager().from(fullPick.getIndices()),
+                            fullPick.getAxis());
+                    return;
                 }
-                PyTorchLibrary.LIB.torchIndexAppendArray(torchIndexHandle, indexArr.getHandle());
-            } else if (elem instanceof NDIndexPick) {
-                // Backward compatible
-                NDIndexFullPick fullPick =
-                        NDIndexFullPick.fromIndex(index, ndArray.getShape()).get();
-                pick(ndArray, ndArray.getManager().from(fullPick.getIndices()), fullPick.getAxis());
-                return;
             }
-        }
-        if (indices.size() == index.getEllipsisIndex()) {
-            PyTorchLibrary.LIB.torchIndexAppendNoneEllipsis(torchIndexHandle, true);
-        }
+            if (indices.size() == index.getEllipsisIndex()) {
+                PyTorchLibrary.LIB.torchIndexAppendNoneEllipsis(torchIndexHandle, true);
+            }
 
-        PyTorchLibrary.LIB.torchIndexAdvPut(
-                ndArray.getHandle(), torchIndexHandle, data.getHandle());
+            PyTorchLibrary.LIB.torchIndexAdvPut(
+                    ndArray.getHandle(), torchIndexHandle, data.getHandle());
+        } finally {
+            PyTorchLibrary.LIB.torchDeleteIndex(torchIndexHandle);
+        }
     }
 
     public static void indexSet(
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 f6c4b3d2c..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);
@@ -214,6 +216,8 @@ final class PyTorchLibrary {
 
     native void torchDeleteTensor(long handle);
 
+    native void torchDeleteIndex(long handle);
+
     native void torchDeleteModule(long handle);
 
     native void torchDeleteIValue(long handle);
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/pytorch/pytorch-native/src/main/native/ai_djl_pytorch_jni_PyTorchLibrary_tensor.cc b/engines/pytorch/pytorch-native/src/main/native/ai_djl_pytorch_jni_PyTorchLibrary_tensor.cc
index f81a66e44..06349f80d 100644
--- a/engines/pytorch/pytorch-native/src/main/native/ai_djl_pytorch_jni_PyTorchLibrary_tensor.cc
+++ b/engines/pytorch/pytorch-native/src/main/native/ai_djl_pytorch_jni_PyTorchLibrary_tensor.cc
@@ -311,6 +311,14 @@ JNIEXPORT void JNICALL Java_ai_djl_pytorch_jni_PyTorchLibrary_torchDeleteTensor(
   API_END()
 }
 
+JNIEXPORT void JNICALL Java_ai_djl_pytorch_jni_PyTorchLibrary_torchDeleteIndex(
+    JNIEnv* env, jobject jthis, jlong jtorch_index_handle) {
+  API_BEGIN()
+  auto* index_ptr = reinterpret_cast<std::vector<torch::indexing::TensorIndex>*>(jtorch_index_handle);
+  delete index_ptr;
+  API_END()
+}
+
 JNIEXPORT jlong JNICALL Java_ai_djl_pytorch_jni_PyTorchLibrary_torchToSparse(
     JNIEnv* env, jobject jthis, jlong jhandle) {
   API_BEGIN()
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");
-    }
 }
diff --git a/examples/build.gradle b/examples/build.gradle
index 1489919fe..8c42e9649 100644
--- a/examples/build.gradle
+++ b/examples/build.gradle
@@ -9,6 +9,7 @@ dependencies {
     implementation project(":model-zoo")
     implementation project(":extensions:timeseries")
     implementation project(":extensions:tokenizers")
+    implementation project(":extensions:audio")
 
     runtimeOnly project(":engines:pytorch:pytorch-model-zoo")
     runtimeOnly project(":engines:tensorflow:tensorflow-model-zoo")
diff --git a/examples/docs/whisper_speech_text.md b/examples/docs/whisper_speech_text.md
new file mode 100644
index 000000000..245a5b8b1
--- /dev/null
+++ b/examples/docs/whisper_speech_text.md
@@ -0,0 +1,57 @@
+# OpenAI Whipser model in DJL
+
+[Whisper](https://github.com/openai/whisper) is an open source model released by OpenAI.
+It can do speech recognition and also machine translation within a single model. In this tutorial,
+we just convert the English portion of the model into Java.
+
+Our input is an audio file:
+
+```
+https://resources.djl.ai/audios/jfk.flac
+```
+
+## Run the example
+
+```
+cd examples
+./gradlew run -Dmain=ai.djl.examples.inference.whisper.SpeechToTextGeneration
+```
+
+Output:
+
+```
+
 
 
 
 And so my fellow Americans , ask not what your country can do for you , ask what you can do for your country . 

+```
+
+## Trace the model
+
+You can use the following script to get the traced model and also vocabulary files:
+
+```python
+from transformers import WhisperProcessor, WhisperForConditionalGeneration
+from datasets import load_dataset
+import torch
+import numpy as np
+
+processor = WhisperProcessor.from_pretrained("openai/whisper-base")
+processor.tokenizer.save_pretrained("whisper-tokenizer")
+model = WhisperForConditionalGeneration.from_pretrained("openai/whisper-base", return_dict=False)
+ds = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation")
+
+test = []
+for ele in ds:
+    test.append(ele["audio"]["array"])
+
+input_features = processor(np.concatenate(test), return_tensors="pt").input_features
+generated_ids = model.generate(inputs=input_features)
+transcription = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
+print("Original: " + transcription)
+
+# Start tracing
+traced_model = torch.jit.trace_module(model, {"generate": [input_features]})
+generated_ids = traced_model.generate(input_features)
+transcription = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
+print("Traced: " + transcription)
+
+torch.jit.save(traced_model, "whisper_en.pt")
+```
\ No newline at end of file
diff --git a/examples/src/main/java/ai/djl/examples/inference/SpeechRecognition.java b/examples/src/main/java/ai/djl/examples/inference/SpeechRecognition.java
index 1561bab1b..5d9a74fad 100644
--- a/examples/src/main/java/ai/djl/examples/inference/SpeechRecognition.java
+++ b/examples/src/main/java/ai/djl/examples/inference/SpeechRecognition.java
@@ -47,8 +47,7 @@ public final class SpeechRecognition {
         logger.info("Result: {}", predict());
     }
 
-    public static String predict()
-            throws UnsupportedAudioFileException, IOException, ModelException, TranslateException {
+    public static String predict() throws IOException, ModelException, TranslateException {
         // Load model.
         // Wav2Vec2 model is a speech model that accepts a float array corresponding to the raw
         // waveform of the speech signal.
@@ -64,7 +63,7 @@ public final class SpeechRecognition {
 
         // Read in audio file
         String wave = "https://resources.djl.ai/audios/speech.wav";
-        Audio audio = AudioFactory.getInstance().fromUrl(wave);
+        Audio audio = AudioFactory.newInstance().fromUrl(wave);
         try (ZooModel<Audio, String> model = criteria.loadModel();
                 Predictor<Audio, String> predictor = model.newPredictor()) {
             return predictor.predict(audio);
diff --git a/examples/src/main/java/ai/djl/examples/inference/whisper/SpeechToTextGeneration.java b/examples/src/main/java/ai/djl/examples/inference/whisper/SpeechToTextGeneration.java
new file mode 100644
index 000000000..de702cd5a
--- /dev/null
+++ b/examples/src/main/java/ai/djl/examples/inference/whisper/SpeechToTextGeneration.java
@@ -0,0 +1,40 @@
+/*
+ * 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.examples.inference.whisper;
+
+import ai.djl.ModelException;
+import ai.djl.training.util.DownloadUtils;
+import ai.djl.translate.TranslateException;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.io.IOException;
+import java.nio.file.Paths;
+
+public final class SpeechToTextGeneration {
+
+    private static final Logger logger = LoggerFactory.getLogger(SpeechToTextGeneration.class);
+
+    private SpeechToTextGeneration() {}
+
+    public static void main(String[] args)
+            throws ModelException, IOException, TranslateException, InterruptedException {
+        DownloadUtils.download(
+                "https://resources.djl.ai/audios/jfk.flac", "build/example/jfk.flac");
+
+        try (WhisperModel model = new WhisperModel()) {
+            logger.info(model.speechToText(Paths.get("build/example/jfk.flac")));
+        }
+    }
+}
diff --git a/examples/src/main/java/ai/djl/examples/inference/whisper/WhisperModel.java b/examples/src/main/java/ai/djl/examples/inference/whisper/WhisperModel.java
new file mode 100644
index 000000000..2cb7cb7a8
--- /dev/null
+++ b/examples/src/main/java/ai/djl/examples/inference/whisper/WhisperModel.java
@@ -0,0 +1,66 @@
+/*
+ * 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.examples.inference.whisper;
+
+import ai.djl.ModelException;
+import ai.djl.inference.Predictor;
+import ai.djl.modality.audio.Audio;
+import ai.djl.modality.audio.AudioFactory;
+import ai.djl.repository.zoo.Criteria;
+import ai.djl.repository.zoo.ZooModel;
+import ai.djl.translate.TranslateException;
+
+import org.bytedeco.ffmpeg.global.avutil;
+
+import java.io.IOException;
+import java.nio.file.Path;
+
+/** An example implementation of OpenAI Whisper Model. */
+public class WhisperModel implements AutoCloseable {
+
+    ZooModel<Audio, String> whisperModel;
+
+    public WhisperModel() throws ModelException, IOException {
+        Criteria<Audio, String> criteria =
+                Criteria.builder()
+                        .setTypes(Audio.class, String.class)
+                        .optModelUrls(
+                                "https://resources.djl.ai/demo/pytorch/whisper/whisper_en.zip")
+                        .optEngine("PyTorch")
+                        .optTranslator(new WhisperTranslator())
+                        .build();
+        whisperModel = criteria.loadModel();
+    }
+
+    public String speechToText(Audio speech) throws TranslateException {
+        try (Predictor<Audio, String> predictor = whisperModel.newPredictor()) {
+            return predictor.predict(speech);
+        }
+    }
+
+    public String speechToText(Path file) throws IOException, TranslateException {
+        Audio audio =
+                AudioFactory.newInstance()
+                        .setChannels(1)
+                        .setSampleRate(16000)
+                        .setSampleFormat(avutil.AV_SAMPLE_FMT_S16)
+                        .fromFile(file);
+        return speechToText(audio);
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public void close() {
+        whisperModel.close();
+    }
+}
diff --git a/examples/src/main/java/ai/djl/examples/inference/whisper/WhisperTranslator.java b/examples/src/main/java/ai/djl/examples/inference/whisper/WhisperTranslator.java
new file mode 100644
index 000000000..155cfc4af
--- /dev/null
+++ b/examples/src/main/java/ai/djl/examples/inference/whisper/WhisperTranslator.java
@@ -0,0 +1,101 @@
+/*
+ * 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.examples.inference.whisper;
+
+import ai.djl.audio.processor.AudioProcessor;
+import ai.djl.audio.processor.LogMelSpectrogram;
+import ai.djl.audio.processor.PadOrTrim;
+import ai.djl.modality.audio.Audio;
+import ai.djl.modality.nlp.DefaultVocabulary;
+import ai.djl.modality.nlp.Vocabulary;
+import ai.djl.ndarray.NDArray;
+import ai.djl.ndarray.NDList;
+import ai.djl.ndarray.NDManager;
+import ai.djl.translate.NoBatchifyTranslator;
+import ai.djl.translate.TranslatorContext;
+import ai.djl.util.JsonUtils;
+
+import com.google.gson.reflect.TypeToken;
+
+import java.io.IOException;
+import java.io.Reader;
+import java.lang.reflect.Type;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.List;
+import java.util.Map;
+
+public class WhisperTranslator implements NoBatchifyTranslator<Audio, String> {
+
+    private List<AudioProcessor> processors;
+    private Vocabulary vocabulary;
+
+    public WhisperTranslator() {
+        processors = new ArrayList<>();
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public void prepare(TranslatorContext ctx) throws IOException {
+        NDManager manager = ctx.getNDManager();
+        Path path = ctx.getModel().getModelPath();
+        Path melFile = path.resolve("mel_80_filters.npz");
+
+        processors.add(new PadOrTrim(480000));
+        processors.add(LogMelSpectrogram.newInstance(melFile, 80, manager));
+
+        Map<String, Integer> vocab;
+        Map<String, Integer> added;
+        Type type = new TypeToken<Map<String, Integer>>() {}.getType();
+        try (Reader reader = Files.newBufferedReader(path.resolve("vocab.json"))) {
+            vocab = JsonUtils.GSON.fromJson(reader, type);
+        }
+        try (Reader reader = Files.newBufferedReader(path.resolve("added_tokens.json"))) {
+            added = JsonUtils.GSON.fromJson(reader, type);
+        }
+        String[] result = new String[vocab.size() + added.size()];
+        vocab.forEach((key, value) -> result[value] = key);
+        added.forEach((key, value) -> result[value] = key);
+        vocabulary = new DefaultVocabulary(Arrays.asList(result));
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public NDList processInput(TranslatorContext ctx, Audio input) throws Exception {
+        NDArray samples = ctx.getNDManager().create(input.getData());
+        for (AudioProcessor processor : processors) {
+            samples = processor.extractFeatures(samples.getManager(), samples);
+        }
+        samples = samples.expandDims(0);
+        NDArray placeholder = ctx.getNDManager().create("");
+        placeholder.setName("module_method:generate");
+        return new NDList(samples, placeholder);
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public String processOutput(TranslatorContext ctx, NDList list) throws Exception {
+        NDArray result = list.singletonOrThrow();
+        List<String> sentence = new ArrayList<>();
+        for (long ele : result.toLongArray()) {
+            sentence.add(vocabulary.getToken(ele));
+            if ("
".equals(vocabulary.getToken(ele))) {
+                break;
+            }
+        }
+        String output = String.join(" ", sentence);
+        return output.replaceAll("[^a-zA-Z0-9<|> ,.!]", "");
+    }
+}
diff --git a/examples/src/main/java/ai/djl/examples/inference/whisper/package-info.java b/examples/src/main/java/ai/djl/examples/inference/whisper/package-info.java
new file mode 100644
index 000000000..8016afbf6
--- /dev/null
+++ b/examples/src/main/java/ai/djl/examples/inference/whisper/package-info.java
@@ -0,0 +1,15 @@
+/*
+ * 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.
+ */
+
+/** Contains examples of Speech to Text task. */
+package ai.djl.examples.inference.whisper;
diff --git a/extensions/audio/src/main/java/ai/djl/audio/FFmpegAudioFactory.java b/extensions/audio/src/main/java/ai/djl/audio/FFmpegAudioFactory.java
index ad6ffa071..07f238708 100644
--- a/extensions/audio/src/main/java/ai/djl/audio/FFmpegAudioFactory.java
+++ b/extensions/audio/src/main/java/ai/djl/audio/FFmpegAudioFactory.java
@@ -14,7 +14,6 @@ package ai.djl.audio;
 
 import ai.djl.modality.audio.Audio;
 import ai.djl.modality.audio.AudioFactory;
-import ai.djl.ndarray.NDArray;
 
 import org.bytedeco.javacv.FFmpegFrameGrabber;
 import org.bytedeco.javacv.Frame;
@@ -22,6 +21,8 @@ import org.bytedeco.javacv.FrameGrabber;
 
 import java.io.IOException;
 import java.io.InputStream;
+import java.nio.Buffer;
+import java.nio.IntBuffer;
 import java.nio.ShortBuffer;
 import java.nio.file.Path;
 import java.util.ArrayList;
@@ -37,6 +38,7 @@ public class FFmpegAudioFactory extends AudioFactory {
     @Override
     public Audio fromFile(Path path) throws IOException {
         try (FFmpegFrameGrabber grabber = new FFmpegFrameGrabber(path.toFile())) {
+            applyConfig(grabber);
             grabber.start();
             float[] floats = grab(grabber);
             return new Audio(floats, grabber.getSampleRate(), grabber.getAudioChannels());
@@ -49,6 +51,7 @@ public class FFmpegAudioFactory extends AudioFactory {
     @Override
     public Audio fromInputStream(InputStream is) throws IOException {
         try (FFmpegFrameGrabber grabber = new FFmpegFrameGrabber(is)) {
+            applyConfig(grabber);
             grabber.start();
             float[] floats = grab(grabber);
             return new Audio(floats, grabber.getSampleRate(), grabber.getAudioChannels());
@@ -57,20 +60,22 @@ public class FFmpegAudioFactory extends AudioFactory {
         }
     }
 
-    /** {@inheritDoc} */
-    @Override
-    public Audio fromData(float[] data) {
-        return new Audio(data);
-    }
-
-    /** {@inheritDoc} */
-    @Override
-    public Audio fromNDArray(NDArray array) {
-        throw new UnsupportedOperationException("Not supported!");
+    private void applyConfig(FFmpegFrameGrabber grabber) {
+        if (channels > 0) {
+            grabber.setAudioChannels(channels);
+        }
+        if (sampleRate > 0) {
+            grabber.setSampleRate(sampleRate);
+        }
+        if (sampleFormat > 0) {
+            grabber.setSampleFormat(sampleFormat);
+        }
     }
 
     /**
-     * Grab frames from the audio using {@link FFmpegFrameGrabber}.
+     * Grabs frames from the audio using {@link FFmpegFrameGrabber}.
+     *
+     * <p>The default channel to grab is 0.
      *
      * @param grabber the {@link FFmpegFrameGrabber}.
      * @return the float array read from the audio.
@@ -79,10 +84,21 @@ public class FFmpegAudioFactory extends AudioFactory {
     private float[] grab(FFmpegFrameGrabber grabber) throws FFmpegFrameGrabber.Exception {
         List<Float> list = new ArrayList<>();
         Frame frame;
-        while ((frame = grabber.grab()) != null) {
-            ShortBuffer buffer = (ShortBuffer) frame.samples[0];
-            for (int i = 0; i < buffer.limit(); i++) {
-                list.add(buffer.get() / (float) Short.MAX_VALUE);
+        while ((frame = grabber.grabFrame(true, false, true, false, false)) != null) {
+            Buffer buf = frame.samples[0];
+            if (buf instanceof ShortBuffer) {
+                ShortBuffer buffer = (ShortBuffer) buf;
+                for (int i = 0; i < buffer.limit(); i++) {
+                    list.add(buffer.get() / (float) Short.MAX_VALUE);
+                }
+            } else if (buf instanceof IntBuffer) {
+                IntBuffer buffer = (IntBuffer) buf;
+                for (int i = 0; i < buffer.limit(); i++) {
+                    list.add(buffer.get() / (float) Integer.MAX_VALUE);
+                }
+            } else {
+                throw new UnsupportedOperationException(
+                        "Unsupported sample format: " + sampleFormat);
             }
         }
         float[] ret = new float[list.size()];
diff --git a/extensions/audio/src/main/java/ai/djl/audio/processor/LogMelSpectrogram.java b/extensions/audio/src/main/java/ai/djl/audio/processor/LogMelSpectrogram.java
new file mode 100644
index 000000000..b76598a45
--- /dev/null
+++ b/extensions/audio/src/main/java/ai/djl/audio/processor/LogMelSpectrogram.java
@@ -0,0 +1,87 @@
+/*
+ * 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.audio.processor;
+
+import ai.djl.ndarray.NDArray;
+import ai.djl.ndarray.NDList;
+import ai.djl.ndarray.NDManager;
+
+import java.io.IOException;
+import java.io.InputStream;
+import java.nio.file.Files;
+import java.nio.file.Path;
+
+/** Apply Log Mel spectrogram to the given data. */
+public class LogMelSpectrogram implements AudioProcessor {
+
+    private static final int N_FFT = 400;
+    private static final int HOP_LENGTH = 160;
+
+    private NDArray melFilters;
+
+    /**
+     * Constructs a new instance of {@code LogMelSpectrogram}.
+     *
+     * @param melFilter the mel filter
+     */
+    public LogMelSpectrogram(NDArray melFilter) {
+        this.melFilters = melFilter;
+    }
+
+    /**
+     * Loads the mel filterbank matrix for projecting STFT into a Mel spectrogram.
+     *
+     * <p>Allows decoupling librosa dependency; saved using: np.savez_compressed( "mel_filters.npz",
+     * mel_80=librosa.filters.mel(sr=16000, n_fft=400, n_mels=80), )
+     *
+     * @param melFile the mdel file saved in .npz format
+     * @param numMel number of mel
+     * @param manager manager to set for the content
+     * @return a new instance of {@code LogMelSpectrogram}
+     * @throws IOException file not loadable
+     */
+    public static LogMelSpectrogram newInstance(Path melFile, int numMel, NDManager manager)
+            throws IOException {
+        try (InputStream is = Files.newInputStream(melFile)) {
+            return newInstance(is, numMel, manager);
+        }
+    }
+
+    /**
+     * Loads the mel filterbank matrix for projecting STFT into a Mel spectrogram.
+     *
+     * @param is the input stream
+     * @param numMel number of mel
+     * @param manager manager to set for the content
+     * @return a new instance of {@code LogMelSpectrogram}
+     * @throws IOException file not loadable
+     */
+    public static LogMelSpectrogram newInstance(InputStream is, int numMel, NDManager manager)
+            throws IOException {
+        return new LogMelSpectrogram(NDList.decode(manager, is).get("mel_" + numMel));
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public NDArray extractFeatures(NDManager manager, NDArray samples) {
+        NDArray window = manager.hanningWindow(N_FFT);
+        NDArray stft = samples.stft(N_FFT, HOP_LENGTH, true, window, true);
+        NDArray magnitudes = stft.get(":,:-1").abs().pow(2);
+        NDArray melSpec = melFilters.matMul(magnitudes);
+        melSpec.attach(manager);
+        NDArray logSpec = melSpec.clip(1e-10, Float.MAX_VALUE).log10();
+        logSpec = logSpec.maximum(logSpec.max().sub(8.0f));
+        logSpec = logSpec.add(4.0f).div(4.0f);
+        return logSpec;
+    }
+}
diff --git a/extensions/audio/src/main/java/ai/djl/audio/processor/PadOrTrim.java b/extensions/audio/src/main/java/ai/djl/audio/processor/PadOrTrim.java
new file mode 100644
index 000000000..368155e3b
--- /dev/null
+++ b/extensions/audio/src/main/java/ai/djl/audio/processor/PadOrTrim.java
@@ -0,0 +1,45 @@
+/*
+ * 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.audio.processor;
+
+import ai.djl.ndarray.NDArray;
+import ai.djl.ndarray.NDManager;
+import ai.djl.ndarray.types.Shape;
+
+/** Pad or trim the samples to the desired numbers. */
+public class PadOrTrim implements AudioProcessor {
+
+    private int desiredSamples;
+
+    /**
+     * Pad or trim the samples to the fixed length.
+     *
+     * @param desiredSamples the desired sample points
+     */
+    public PadOrTrim(int desiredSamples) {
+        this.desiredSamples = desiredSamples;
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public NDArray extractFeatures(NDManager manager, NDArray samples) {
+        if (samples.getShape().dimension() != 1) {
+            throw new UnsupportedOperationException("Batch samples not supported.");
+        }
+        long sampleLength = samples.getShape().get(0);
+        if (sampleLength > desiredSamples) {
+            return samples.get(":" + desiredSamples);
+        }
+        return samples.concat(manager.zeros(new Shape(desiredSamples - sampleLength)));
+    }
+}
diff --git a/gradle/wrapper/gradle-wrapper.properties b/gradle/wrapper/gradle-wrapper.properties
index 691737c09..1345e3c22 100644
--- a/gradle/wrapper/gradle-wrapper.properties
+++ b/gradle/wrapper/gradle-wrapper.properties
@@ -3,4 +3,4 @@ distributionBase=GRADLE_USER_HOME
 distributionPath=wrapper/dists
 zipStoreBase=GRADLE_USER_HOME
 zipStorePath=wrapper/dists
-distributionUrl=https\://services.gradle.org/distributions/gradle-7.5.1-bin.zip
+distributionUrl=https\://services.gradle.org/distributions/gradle-7.6-bin.zip
diff --git a/integration/build.gradle b/integration/build.gradle
index 675d59061..798d8165e 100644
--- a/integration/build.gradle
+++ b/integration/build.gradle
@@ -52,33 +52,6 @@ run {
     jvmArgs "-Xverify:none"
 }
 
-jacocoTestReport {
-    group = "Reporting"
-
-    javaProjects().each { p ->
-        if ([":api", ":basicdataset"].contains(p.getPath())) {
-            additionalSourceDirs files((Set<File>) p.sourceSets.main.allJava.srcDirs)
-            additionalClassDirs((FileCollection) p.sourceSets.main.output)
-        }
-    }
-
-    afterEvaluate {
-        getClassDirectories().setFrom(classDirectories.files.collect {
-            fileTree(dir: it, exclude: ["**/ai/djl/integration/**"])
-        })
-    }
-
-    reports {
-        xml.required = true
-        csv.required = false
-    }
-
-    // Work-around to allow us to build list of executionData files in doFirst
-    onlyIf = {
-        true
-    }
-}
-
 task debugEnv(type: JavaExec) {
     classpath = sourceSets.main.runtimeClasspath
     systemProperties System.getProperties()
@@ -88,32 +61,3 @@ task debugEnv(type: JavaExec) {
 }
 
 tasks.distTar.enabled = false
-test.finalizedBy project.tasks.jacocoTestReport, project.tasks.jacocoTestCoverageVerification
-
-jacocoTestCoverageVerification {
-    javaProjects().each { p ->
-        if ([":api", ":basicdataset"].contains(p.getPath())) {
-            additionalSourceDirs files((Set<File>) p.sourceSets.main.allJava.srcDirs)
-            additionalClassDirs((FileCollection) p.sourceSets.main.output)
-        }
-    }
-
-    afterEvaluate {
-        getClassDirectories().setFrom(classDirectories.files.collect {
-            fileTree(dir: it, exclude: ["**/ai/djl/integration/**"])
-        })
-    }
-
-    violationRules {
-        rule {
-            limit {
-                if (System.getProperty("ai.djl.default_engine") != null) {
-                    minimum = 0.2
-                } else {
-                    // TODO: improve coverage by running integration tests with 3 engines
-                    minimum = 0.4
-                }
-            }
-        }
-    }
-}
diff --git a/integration/src/main/java/ai/djl/integration/util/DebugEnvironment.java b/integration/src/main/java/ai/djl/integration/util/DebugEnvironment.java
index f3146c6f2..cd352ca7d 100644
--- a/integration/src/main/java/ai/djl/integration/util/DebugEnvironment.java
+++ b/integration/src/main/java/ai/djl/integration/util/DebugEnvironment.java
@@ -73,7 +73,8 @@ public final class DebugEnvironment {
 
         if (!TestUtils.isWindows()) {
             System.out.println("GCC: ");
-            Process process = rt.exec("gcc --version");
+            String[] cmd = {"gcc", "--version"};
+            Process process = rt.exec(cmd);
             try (Scanner gccOut =
                     new Scanner(process.getInputStream(), StandardCharsets.UTF_8.name())) {
                 gccOut.useDelimiter(System.lineSeparator());
diff --git a/jacoco/gradlew b/jacoco/gradlew
new file mode 120000
index 000000000..502f5a2d3
--- /dev/null
+++ b/jacoco/gradlew
@@ -0,0 +1,1 @@
+../gradlew
\ No newline at end of file


Test patch:
diff --git a/api/src/test/java/ai/djl/modality/audio/AudioFactoryTest.java b/api/src/test/java/ai/djl/modality/audio/AudioFactoryTest.java
index 40ec26d27..c5b082c7a 100644
--- a/api/src/test/java/ai/djl/modality/audio/AudioFactoryTest.java
+++ b/api/src/test/java/ai/djl/modality/audio/AudioFactoryTest.java
@@ -22,8 +22,6 @@ import org.testng.annotations.BeforeClass;
 import org.testng.annotations.Test;
 
 import java.io.IOException;
-import java.io.InputStream;
-import java.nio.file.Files;
 import java.nio.file.Paths;
 
 public class AudioFactoryTest {
@@ -37,37 +35,36 @@ public class AudioFactoryTest {
 
     @Test
     public void testFromFile() throws IOException {
-        Audio audio = AudioFactory.getInstance().fromFile(Paths.get("build/test/test_01.wav"));
+        Audio audio = AudioFactory.newInstance().fromFile(Paths.get("build/test/test_01.wav"));
         Assert.assertEquals(audio.getSampleRate(), 16000f);
         Assert.assertEquals(audio.getChannels(), 1);
+
+        AudioFactory factory = AudioFactory.newInstance().setChannels(0);
+        Assert.assertEquals(factory.getChannels(), 0);
+        Assert.assertEquals(factory.getSampleFormat(), 0);
+        Assert.assertEquals(factory.getSampleRate(), 0);
+        Assert.assertThrows(() -> factory.setChannels(1));
+        Assert.assertThrows(() -> factory.setSampleRate(1000));
+        Assert.assertThrows(() -> factory.setSampleFormat(1)); // AV_SAMPLE_FMT_S16
     }
 
     @Test
     public void testFromUrl() throws IOException {
-        Audio audio = AudioFactory.getInstance().fromUrl("build/test/test_01.wav");
+        Audio audio = AudioFactory.newInstance().fromUrl("build/test/test_01.wav");
         Assert.assertEquals(audio.getSampleRate(), 16000f);
         Assert.assertEquals(audio.getChannels(), 1);
 
-        audio = AudioFactory.getInstance().fromUrl(URL);
+        audio = AudioFactory.newInstance().fromUrl(URL);
         Assert.assertEquals(audio.getSampleRate(), 16000f);
         Assert.assertEquals(audio.getChannels(), 1);
     }
 
-    @Test
-    public void testFromInputStream() throws IOException {
-        try (InputStream is = Files.newInputStream(Paths.get("build/test/test_01.wav"))) {
-            Audio audio = AudioFactory.getInstance().fromInputStream(is);
-            Assert.assertEquals(audio.getSampleRate(), 16000f);
-            Assert.assertEquals(audio.getChannels(), 1);
-        }
-    }
-
     @Test
     public void testFromData() {
         float[] data = {0.001f, 0.002f, 0.003f};
-        Audio audio = AudioFactory.getInstance().fromData(data);
+        Audio audio = AudioFactory.newInstance().fromData(data);
         Assert.assertEquals(audio.getData(), data);
-        Assert.assertEquals(audio.getSampleRate(), 0);
+        Assert.assertEquals(audio.getSampleRate(), 0f);
         Assert.assertEquals(audio.getChannels(), 0);
     }
 
@@ -75,7 +72,7 @@ public class AudioFactoryTest {
     public void testFromNDArray() {
         try (NDManager manager = NDManager.newBaseManager()) {
             NDArray array = manager.zeros(new Shape(1));
-            AudioFactory.getInstance().fromNDArray(array);
+            AudioFactory.newInstance().fromNDArray(array);
         }
     }
 }
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/extensions/audio/src/test/java/ai/djl/audio/FFmpegAudioFactoryTest.java b/extensions/audio/src/test/java/ai/djl/audio/FFmpegAudioFactoryTest.java
index 7cfe2ca63..43da44d8b 100644
--- a/extensions/audio/src/test/java/ai/djl/audio/FFmpegAudioFactoryTest.java
+++ b/extensions/audio/src/test/java/ai/djl/audio/FFmpegAudioFactoryTest.java
@@ -14,53 +14,56 @@ package ai.djl.audio;
 
 import ai.djl.modality.audio.Audio;
 import ai.djl.modality.audio.AudioFactory;
-import ai.djl.ndarray.NDArray;
-import ai.djl.ndarray.NDManager;
-import ai.djl.ndarray.types.Shape;
 import ai.djl.training.util.DownloadUtils;
 
+import org.bytedeco.ffmpeg.global.avutil;
 import org.testng.Assert;
 import org.testng.annotations.BeforeClass;
 import org.testng.annotations.Test;
 
 import java.io.IOException;
+import java.net.URL;
 import java.nio.file.Paths;
 
 public class FFmpegAudioFactoryTest {
 
-    private static final String URL = "https://resources.djl.ai/audios/test_01.wav";
-
     @BeforeClass
     public void setUp() throws IOException {
-        DownloadUtils.download(URL, "build/test/test_01.wav");
+        DownloadUtils.download(
+                "https://resources.djl.ai/audios/test_01.wav", "build/test/test_01.wav");
     }
 
     @Test
     public void testFactory() throws IOException {
-        Audio audio = AudioFactory.getInstance().fromFile(Paths.get("build/test/test_01.wav"));
-        Assert.assertEquals(audio.getSampleRate(), 16000f);
-        Assert.assertEquals(audio.getChannels(), 1);
+        AudioFactory factory =
+                AudioFactory.newInstance()
+                        .setChannels(1)
+                        .setSampleRate(10000)
+                        .setSampleFormat(avutil.AV_SAMPLE_FMT_S16);
 
-        audio = AudioFactory.getInstance().fromUrl("build/test/test_01.wav");
-        Assert.assertEquals(audio.getSampleRate(), 16000f);
-        Assert.assertEquals(audio.getChannels(), 1);
+        Assert.assertEquals(factory.getChannels(), 1);
+        Assert.assertEquals(factory.getSampleRate(), 10000);
+        Assert.assertEquals(factory.getSampleFormat(), avutil.AV_SAMPLE_FMT_S16);
 
-        audio = AudioFactory.getInstance().fromUrl(URL);
-        Assert.assertEquals(audio.getSampleRate(), 16000f);
+        Audio audio = factory.fromUrl("build/test/test_01.wav");
+        Assert.assertEquals(audio.getSampleRate(), 10000f);
         Assert.assertEquals(audio.getChannels(), 1);
 
-        float[] data = {0.001f, 0.002f, 0.003f};
-        audio = AudioFactory.getInstance().fromData(data);
-        Assert.assertEquals(audio.getData(), data);
-        Assert.assertEquals(audio.getSampleRate(), 0);
-        Assert.assertEquals(audio.getChannels(), 0);
-    }
+        URL url = Paths.get("build/test/test_01.wav").toAbsolutePath().toUri().toURL();
+
+        audio =
+                AudioFactory.newInstance()
+                        .setSampleRate(8000)
+                        .setSampleFormat(avutil.AV_SAMPLE_FMT_S32)
+                        .fromUrl(url);
+        Assert.assertEquals(audio.getSampleRate(), 8000f);
+        Assert.assertEquals(audio.getChannels(), 1);
 
-    @Test(expectedExceptions = UnsupportedOperationException.class)
-    public void testFromNDArray() {
-        try (NDManager manager = NDManager.newBaseManager()) {
-            NDArray array = manager.zeros(new Shape(1));
-            AudioFactory.getInstance().fromNDArray(array);
-        }
+        Assert.assertThrows(
+                () -> {
+                    AudioFactory.newInstance()
+                            .setSampleFormat(avutil.AV_SAMPLE_FMT_DBL)
+                            .fromUrl(url);
+                });
     }
 }
diff --git a/extensions/audio/src/test/java/ai/djl/audio/processor/AudioProcessorTest.java b/extensions/audio/src/test/java/ai/djl/audio/processor/AudioProcessorTest.java
index 8070a1842..4b9517369 100644
--- a/extensions/audio/src/test/java/ai/djl/audio/processor/AudioProcessorTest.java
+++ b/extensions/audio/src/test/java/ai/djl/audio/processor/AudioProcessorTest.java
@@ -26,16 +26,19 @@ import org.testng.annotations.BeforeClass;
 import org.testng.annotations.Test;
 
 import java.io.IOException;
+import java.nio.file.Paths;
 import java.util.Arrays;
 import java.util.Collections;
 
 public class AudioProcessorTest {
 
-    private static final String URL = "https://resources.djl.ai/audios/test_01.wav";
+    private static final String URL = "https://resources.djl.ai/";
 
     @BeforeClass
     public void setUp() throws IOException {
-        DownloadUtils.download(URL, "build/test/test_01.wav");
+        DownloadUtils.download(URL + "audios/test_01.wav", "build/test/test_01.wav");
+        DownloadUtils.download(
+                URL + "demo/pytorch/whisper/mel_80_filters.npz", "build/test/mel_80_filters.npz");
     }
 
     @Test
@@ -50,6 +53,31 @@ public class AudioProcessorTest {
         Assert.assertEquals(AudioUtils.rmsDb(samples), -20.0f, 1e-3);
     }
 
+    @Test
+    public void testPadOrTrim() {
+        try (NDManager manager = NDManager.newBaseManager()) {
+            AudioProcessor processor = new PadOrTrim(100);
+            NDArray test = manager.ones(new Shape(10));
+            NDArray result = processor.extractFeatures(manager, test);
+            Assert.assertEquals(result.getShape(), new Shape(100));
+            test = manager.ones(new Shape(1000));
+            result = processor.extractFeatures(manager, test);
+            Assert.assertEquals(result.getShape(), new Shape(100));
+        }
+    }
+
+    @Test
+    void testLogMelSpectrogram() throws IOException {
+        try (NDManager manager = NDManager.newBaseManager()) {
+            AudioProcessor processor =
+                    LogMelSpectrogram.newInstance(
+                            Paths.get("build/test/mel_80_filters.npz"), 80, manager);
+            NDArray samples = manager.zeros(new Shape(480000));
+            NDArray result = processor.extractFeatures(manager, samples);
+            Assert.assertEquals(result.getShape(), new Shape(80, 3000));
+        }
+    }
+
     @Test
     public static void testLinearSpecgram() {
         NDManager manager = NDManager.newBaseManager(Device.cpu());
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())) {