Repository: deepjavalibrary/djl
Pull request #1346 (2021-11-04 22:10:46)

Predictor with device
Adds a new API to create a predictor with a specified device. This makes it easier to create multiple predictors with different devices from the same model to support use cases like multiple GPU.

codecov-commenter: # [Codecov](https://codecov.io/gh/deepjavalibrary/djl/pull/1346?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) Report
> Merging [#1346](https://codecov.io/gh/deepjavalibrary/djl/pull/1346?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) (f24c5f0) into [master](https://codecov.io/gh/deepjavalibrary/djl/commit/bb5073f7ffc33f1883621fc1783dbb01e70b418f?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) (bb5073f) will **decrease** coverage by `0.37%`.
> The diff coverage is `55.10%`.

[![Impacted file tree graph](https://codecov.io/gh/deepjavalibrary/djl/pull/1346/graphs/tree.svg?width=650&height=150&src=pr&token=mDKVNwFXRV&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary)](https://codecov.io/gh/deepjavalibrary/djl/pull/1346?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary)

```diff
@@             Coverage Diff              @@
##             master    #1346      +/-   ##
============================================
- Coverage     72.08%   71.71%   -0.38%     
- Complexity     5126     5180      +54     
============================================
  Files           473      480       +7     
  Lines         21970    22453     +483     
  Branches       2351     2411      +60     
============================================
+ Hits          15838    16102     +264     
- Misses         4925     5129     +204     
- Partials       1207     1222      +15     
```


| [Impacted Files](https://codecov.io/gh/deepjavalibrary/djl/pull/1346?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) | Coverage Δ | |
|---|---|---|
| [.../java/ai/djl/modality/cv/BufferedImageFactory.java](https://codecov.io/gh/deepjavalibrary/djl/pull/1346/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L0J1ZmZlcmVkSW1hZ2VGYWN0b3J5LmphdmE=) | `72.72% <ø> (+0.27%)` | :arrow_up: |
| [...i/djl/modality/cv/translator/BigGANTranslator.java](https://codecov.io/gh/deepjavalibrary/djl/pull/1346/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3RyYW5zbGF0b3IvQmlnR0FOVHJhbnNsYXRvci5qYXZh) | `21.42% <ø> (-5.24%)` | :arrow_down: |
| [...odality/cv/translator/BigGANTranslatorFactory.java](https://codecov.io/gh/deepjavalibrary/djl/pull/1346/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3RyYW5zbGF0b3IvQmlnR0FOVHJhbnNsYXRvckZhY3RvcnkuamF2YQ==) | `33.33% <0.00%> (+8.33%)` | :arrow_up: |
| [...nslator/InstanceSegmentationTranslatorFactory.java](https://codecov.io/gh/deepjavalibrary/djl/pull/1346/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3RyYW5zbGF0b3IvSW5zdGFuY2VTZWdtZW50YXRpb25UcmFuc2xhdG9yRmFjdG9yeS5qYXZh) | `14.28% <0.00%> (-3.90%)` | :arrow_down: |
| [.../modality/cv/translator/YoloTranslatorFactory.java](https://codecov.io/gh/deepjavalibrary/djl/pull/1346/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3RyYW5zbGF0b3IvWW9sb1RyYW5zbGF0b3JGYWN0b3J5LmphdmE=) | `8.33% <0.00%> (-1.67%)` | :arrow_down: |
| [...i/djl/modality/cv/translator/YoloV5Translator.java](https://codecov.io/gh/deepjavalibrary/djl/pull/1346/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3RyYW5zbGF0b3IvWW9sb1Y1VHJhbnNsYXRvci5qYXZh) | `5.69% <0.00%> (ø)` | |
| [...odality/cv/translator/YoloV5TranslatorFactory.java](https://codecov.io/gh/deepjavalibrary/djl/pull/1346/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3RyYW5zbGF0b3IvWW9sb1Y1VHJhbnNsYXRvckZhY3RvcnkuamF2YQ==) | `8.33% <0.00%> (-1.67%)` | :arrow_down: |
| [...modality/nlp/embedding/TrainableWordEmbedding.java](https://codecov.io/gh/deepjavalibrary/djl/pull/1346/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L25scC9lbWJlZGRpbmcvVHJhaW5hYmxlV29yZEVtYmVkZGluZy5qYXZh) | `45.00% <0.00%> (-3.34%)` | :arrow_down: |
| [...pi/src/main/java/ai/djl/ndarray/BytesSupplier.java](https://codecov.io/gh/deepjavalibrary/djl/pull/1346/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL25kYXJyYXkvQnl0ZXNTdXBwbGllci5qYXZh) | `54.54% <0.00%> (-12.13%)` | :arrow_down: |
| [...i/src/main/java/ai/djl/ndarray/NDArrayAdapter.java](https://codecov.io/gh/deepjavalibrary/djl/pull/1346/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL25kYXJyYXkvTkRBcnJheUFkYXB0ZXIuamF2YQ==) | `15.23% <0.00%> (-0.83%)` | :arrow_down: |
| ... and [114 more](https://codecov.io/gh/deepjavalibrary/djl/pull/1346/diff?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) | |

------

[Continue to review full report at Codecov](https://codecov.io/gh/deepjavalibrary/djl/pull/1346?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/deepjavalibrary/djl/pull/1346?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary). Last update [518cadc...f24c5f0](https://codecov.io/gh/deepjavalibrary/djl/pull/1346?src=pr&el=lastupdated&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary).


Patch:
diff --git a/api/src/main/java/ai/djl/BaseModel.java b/api/src/main/java/ai/djl/BaseModel.java
index a27cf5a96..eab4001de 100644
--- a/api/src/main/java/ai/djl/BaseModel.java
+++ b/api/src/main/java/ai/djl/BaseModel.java
@@ -99,8 +99,8 @@ public abstract class BaseModel implements Model {
 
     /** {@inheritDoc} */
     @Override
-    public <I, O> Predictor<I, O> newPredictor(Translator<I, O> translator) {
-        return new Predictor<>(this, translator, false);
+    public <I, O> Predictor<I, O> newPredictor(Translator<I, O> translator, Device device) {
+        return new Predictor<>(this, translator, device, false);
     }
 
     /** {@inheritDoc} */
diff --git a/api/src/main/java/ai/djl/Model.java b/api/src/main/java/ai/djl/Model.java
index 34ccc65d7..fb5af33fb 100644
--- a/api/src/main/java/ai/djl/Model.java
+++ b/api/src/main/java/ai/djl/Model.java
@@ -203,15 +203,28 @@ public interface Model extends AutoCloseable {
      */
     Trainer newTrainer(TrainingConfig trainingConfig);
 
+    /**
+     * Creates a new Predictor based on the model on the current device.
+     *
+     * @param translator the object used for pre-processing and postprocessing
+     * @param <I> the input object for pre-processing
+     * @param <O> the output object from postprocessing
+     * @return an instance of {@code Predictor}
+     */
+    default <I, O> Predictor<I, O> newPredictor(Translator<I, O> translator) {
+        return newPredictor(translator, getNDManager().getDevice());
+    }
+
     /**
      * Creates a new Predictor based on the model.
      *
      * @param translator the object used for pre-processing and postprocessing
+     * @param device the device to use for prediction
      * @param <I> the input object for pre-processing
      * @param <O> the output object from postprocessing
      * @return an instance of {@code Predictor}
      */
-    <I, O> Predictor<I, O> newPredictor(Translator<I, O> translator);
+    <I, O> Predictor<I, O> newPredictor(Translator<I, O> translator, Device device);
 
     /**
      * Returns the input descriptor of the model.
diff --git a/api/src/main/java/ai/djl/inference/Predictor.java b/api/src/main/java/ai/djl/inference/Predictor.java
index edde46239..7808b8cf2 100644
--- a/api/src/main/java/ai/djl/inference/Predictor.java
+++ b/api/src/main/java/ai/djl/inference/Predictor.java
@@ -12,6 +12,7 @@
  */
 package ai.djl.inference;
 
+import ai.djl.Device;
 import ai.djl.Model;
 import ai.djl.metric.Metrics;
 import ai.djl.ndarray.LazyNDArray;
@@ -94,11 +95,17 @@ public class Predictor<I, O> implements AutoCloseable {
      *
      * @param model the model on which the predictions are based
      * @param translator the translator to be used
-     * @param copy whether to copy the parameters to the parameter store
+     * @param device the device for prediction
+     * @param copy whether to copy the parameters to the parameter store. If the device changes, it
+     *     will copy regardless
      */
-    public Predictor(Model model, Translator<I, O> translator, boolean copy) {
+    public Predictor(Model model, Translator<I, O> translator, Device device, boolean copy) {
+        if (!device.equals(model.getNDManager().getDevice())) {
+            // Always copy during device changes
+            copy = true;
+        }
         this.model = model;
-        this.manager = model.getNDManager().newSubManager();
+        this.manager = model.getNDManager().newSubManager(device);
         this.manager.setName("predictor");
         this.translator = translator;
         block = model.getBlock();
diff --git a/api/src/main/java/ai/djl/repository/zoo/ZooModel.java b/api/src/main/java/ai/djl/repository/zoo/ZooModel.java
index 07414620a..e0cef0bbb 100644
--- a/api/src/main/java/ai/djl/repository/zoo/ZooModel.java
+++ b/api/src/main/java/ai/djl/repository/zoo/ZooModel.java
@@ -12,6 +12,7 @@
  */
 package ai.djl.repository.zoo;
 
+import ai.djl.Device;
 import ai.djl.Model;
 import ai.djl.inference.Predictor;
 import ai.djl.ndarray.NDManager;
@@ -124,10 +125,21 @@ public class ZooModel<I, O> implements Model {
         return newPredictor(translator);
     }
 
+    /**
+     * Creates a new Predictor based on the model with the default translator and a specified
+     * device.
+     *
+     * @param device the device to use for prediction
+     * @return an instance of {@code Predictor}
+     */
+    public Predictor<I, O> newPredictor(Device device) {
+        return model.newPredictor(translator, device);
+    }
+
     /** {@inheritDoc} */
     @Override
-    public <P, Q> Predictor<P, Q> newPredictor(Translator<P, Q> translator) {
-        return model.newPredictor(translator);
+    public <P, Q> Predictor<P, Q> newPredictor(Translator<P, Q> translator, Device device) {
+        return model.newPredictor(translator, device);
     }
 
     /**
diff --git a/engines/dlr/dlr-engine/src/main/java/ai/djl/dlr/engine/DlrModel.java b/engines/dlr/dlr-engine/src/main/java/ai/djl/dlr/engine/DlrModel.java
index a0b43ce44..5efb7102f 100644
--- a/engines/dlr/dlr-engine/src/main/java/ai/djl/dlr/engine/DlrModel.java
+++ b/engines/dlr/dlr-engine/src/main/java/ai/djl/dlr/engine/DlrModel.java
@@ -13,6 +13,7 @@
 package ai.djl.dlr.engine;
 
 import ai.djl.BaseModel;
+import ai.djl.Device;
 import ai.djl.Model;
 import ai.djl.inference.Predictor;
 import ai.djl.ndarray.NDManager;
@@ -61,8 +62,8 @@ public class DlrModel extends BaseModel {
 
     /** {@inheritDoc} */
     @Override
-    public <I, O> Predictor<I, O> newPredictor(Translator<I, O> translator) {
-        return new DlrPredictor<>(this, modelDir.toString(), manager.getDevice(), translator);
+    public <I, O> Predictor<I, O> newPredictor(Translator<I, O> translator, Device device) {
+        return new DlrPredictor<>(this, modelDir.toString(), device, translator);
     }
 
     private void checkModelFiles(String prefix) throws IOException {
diff --git a/engines/dlr/dlr-engine/src/main/java/ai/djl/dlr/engine/DlrPredictor.java b/engines/dlr/dlr-engine/src/main/java/ai/djl/dlr/engine/DlrPredictor.java
index b5447066a..4a332fcc2 100644
--- a/engines/dlr/dlr-engine/src/main/java/ai/djl/dlr/engine/DlrPredictor.java
+++ b/engines/dlr/dlr-engine/src/main/java/ai/djl/dlr/engine/DlrPredictor.java
@@ -34,7 +34,7 @@ public class DlrPredictor<I, O> extends Predictor<I, O> {
      */
     public DlrPredictor(
             DlrModel model, String modelDir, Device device, Translator<I, O> translator) {
-        super(model, translator, false);
+        super(model, translator, device, false);
         long modelHandle = JniUtils.createDlrModel(modelDir, device);
         block = new DlrSymbolBlock((DlrNDManager) manager, modelHandle);
         // disable cpu affinity by default
diff --git a/engines/paddlepaddle/paddlepaddle-engine/src/main/java/ai/djl/paddlepaddle/engine/PpModel.java b/engines/paddlepaddle/paddlepaddle-engine/src/main/java/ai/djl/paddlepaddle/engine/PpModel.java
index 974a3e712..0947356a4 100644
--- a/engines/paddlepaddle/paddlepaddle-engine/src/main/java/ai/djl/paddlepaddle/engine/PpModel.java
+++ b/engines/paddlepaddle/paddlepaddle-engine/src/main/java/ai/djl/paddlepaddle/engine/PpModel.java
@@ -124,8 +124,8 @@ public class PpModel extends BaseModel {
 
     /** {@inheritDoc} */
     @Override
-    public <I, O> Predictor<I, O> newPredictor(Translator<I, O> translator) {
-        return new PpPredictor<>(this, paddlePredictor.copy(), translator);
+    public <I, O> Predictor<I, O> newPredictor(Translator<I, O> translator, Device device) {
+        return new PpPredictor<>(this, paddlePredictor.copy(), translator, device);
     }
 
     /** {@inheritDoc} */
diff --git a/engines/paddlepaddle/paddlepaddle-engine/src/main/java/ai/djl/paddlepaddle/engine/PpPredictor.java b/engines/paddlepaddle/paddlepaddle-engine/src/main/java/ai/djl/paddlepaddle/engine/PpPredictor.java
index 74f4559cf..6aaeda0e2 100644
--- a/engines/paddlepaddle/paddlepaddle-engine/src/main/java/ai/djl/paddlepaddle/engine/PpPredictor.java
+++ b/engines/paddlepaddle/paddlepaddle-engine/src/main/java/ai/djl/paddlepaddle/engine/PpPredictor.java
@@ -12,6 +12,7 @@
  */
 package ai.djl.paddlepaddle.engine;
 
+import ai.djl.Device;
 import ai.djl.Model;
 import ai.djl.inference.Predictor;
 import ai.djl.translate.Translator;
@@ -31,9 +32,11 @@ public class PpPredictor<I, O> extends Predictor<I, O> {
      * @param model the model on which the predictions are based
      * @param predictor the C++ Paddle Predictor handle
      * @param translator the translator to be used
+     * @param device the device to be used
      */
-    public PpPredictor(Model model, PaddlePredictor predictor, Translator<I, O> translator) {
-        super(model, translator, false);
+    public PpPredictor(
+            Model model, PaddlePredictor predictor, Translator<I, O> translator, Device device) {
+        super(model, translator, device, false);
         this.predictor = predictor;
         block = new PpSymbolBlock(predictor, (PpNDManager) model.getNDManager());
     }
diff --git a/engines/tensorrt/src/main/java/ai/djl/tensorrt/engine/TrtModel.java b/engines/tensorrt/src/main/java/ai/djl/tensorrt/engine/TrtModel.java
index 5a8cf65ae..219350941 100644
--- a/engines/tensorrt/src/main/java/ai/djl/tensorrt/engine/TrtModel.java
+++ b/engines/tensorrt/src/main/java/ai/djl/tensorrt/engine/TrtModel.java
@@ -13,6 +13,7 @@
 package ai.djl.tensorrt.engine;
 
 import ai.djl.BaseModel;
+import ai.djl.Device;
 import ai.djl.Model;
 import ai.djl.inference.Predictor;
 import ai.djl.ndarray.types.DataType;
@@ -78,10 +79,11 @@ public class TrtModel extends BaseModel {
 
     /** {@inheritDoc} */
     @Override
-    public <I, O> Predictor<I, O> newPredictor(Translator<I, O> translator) {
+    public <I, O> Predictor<I, O> newPredictor(Translator<I, O> translator, Device device) {
         TrtSymbolBlock trtSymbol = ((TrtSymbolBlock) block);
-        TrtSession session = trtSymbol.createSession((TrtNDManager) manager);
-        return new TrtPredictor<>(this, translator, session);
+        TrtNDManager predManager = ((TrtNDManager) manager).newSubManager(device);
+        TrtSession session = trtSymbol.createSession(predManager);
+        return new TrtPredictor<>(this, translator, device, session);
     }
 
     private Path findModelFile(String prefix) {
diff --git a/engines/tensorrt/src/main/java/ai/djl/tensorrt/engine/TrtPredictor.java b/engines/tensorrt/src/main/java/ai/djl/tensorrt/engine/TrtPredictor.java
index 4e61587b7..b096bb7af 100644
--- a/engines/tensorrt/src/main/java/ai/djl/tensorrt/engine/TrtPredictor.java
+++ b/engines/tensorrt/src/main/java/ai/djl/tensorrt/engine/TrtPredictor.java
@@ -12,13 +12,14 @@
  */
 package ai.djl.tensorrt.engine;
 
+import ai.djl.Device;
 import ai.djl.inference.Predictor;
 import ai.djl.translate.Translator;
 
 class TrtPredictor<I, O> extends Predictor<I, O> {
 
-    TrtPredictor(TrtModel model, Translator<I, O> translator, TrtSession session) {
-        super(model, translator, false);
+    TrtPredictor(TrtModel model, Translator<I, O> translator, Device device, TrtSession session) {
+        super(model, translator, device, false);
         block = session;
     }
 
diff --git a/extensions/fasttext/src/main/java/ai/djl/fasttext/FtModel.java b/extensions/fasttext/src/main/java/ai/djl/fasttext/FtModel.java
index f558b7d8f..c473dcfc6 100644
--- a/extensions/fasttext/src/main/java/ai/djl/fasttext/FtModel.java
+++ b/extensions/fasttext/src/main/java/ai/djl/fasttext/FtModel.java
@@ -12,6 +12,7 @@
  */
 package ai.djl.fasttext;
 
+import ai.djl.Device;
 import ai.djl.MalformedModelException;
 import ai.djl.Model;
 import ai.djl.basicdataset.RawDataset;
@@ -203,8 +204,8 @@ public class FtModel implements Model {
 
     /** {@inheritDoc} */
     @Override
-    public <I, O> Predictor<I, O> newPredictor(Translator<I, O> translator) {
-        return new Predictor<>(this, translator, false);
+    public <I, O> Predictor<I, O> newPredictor(Translator<I, O> translator, Device device) {
+        return new Predictor<>(this, translator, device, false);
     }
 
     /** {@inheritDoc} */


Test patch:
diff --git a/integration/src/main/java/ai/djl/integration/tests/inference/PredictorTest.java b/integration/src/main/java/ai/djl/integration/tests/inference/PredictorTest.java
new file mode 100644
index 000000000..66ffda367
--- /dev/null
+++ b/integration/src/main/java/ai/djl/integration/tests/inference/PredictorTest.java
@@ -0,0 +1,82 @@
+/*
+ * Copyright 2019 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.integration.tests.inference;
+
+import ai.djl.Device;
+import ai.djl.Model;
+import ai.djl.basicmodelzoo.basic.Mlp;
+import ai.djl.inference.Predictor;
+import ai.djl.ndarray.NDList;
+import ai.djl.ndarray.NDManager;
+import ai.djl.ndarray.types.Shape;
+import ai.djl.nn.Block;
+import ai.djl.testing.TestRequirements;
+import ai.djl.training.DefaultTrainingConfig;
+import ai.djl.training.Trainer;
+import ai.djl.training.TrainingConfig;
+import ai.djl.training.loss.Loss;
+import ai.djl.translate.NoopTranslator;
+import ai.djl.translate.TranslateException;
+import org.testng.Assert;
+import org.testng.annotations.Test;
+
+public class PredictorTest {
+
+    @Test
+    public void testPredictorWithDeviceGpu() throws TranslateException {
+        TestRequirements.gpu();
+
+        // CPU model, GPU predictor
+        predictWithDeviceHelper(Device.cpu(), Device.gpu());
+    }
+
+    @Test
+    public void testPredictorWithDeviceCpu() throws TranslateException {
+        TestRequirements.gpu();
+
+        // GPU model, CPU predictor
+        predictWithDeviceHelper(Device.gpu(), Device.cpu());
+    }
+
+    public void predictWithDeviceHelper(Device modelDevice, Device predictorDevice)
+            throws TranslateException {
+        // Create simple model on modelDevice
+        try (NDManager manager = NDManager.newBaseManager(modelDevice)) {
+            Block block = new Mlp(10, 10, new int[] {10});
+            Model model = Model.newInstance("mlp");
+            model.setBlock(block);
+
+            // Use trainer for initialization
+            TrainingConfig config =
+                    new DefaultTrainingConfig(Loss.softmaxCrossEntropyLoss())
+                            .optDevices(new Device[] {modelDevice});
+            try (Trainer trainer = model.newTrainer(config)) {
+                trainer.initialize(new Shape(1, 10));
+            }
+
+            // Create predictor on predictorDevice
+            try (Predictor<NDList, NDList> predictor =
+                    model.newPredictor(new NoopTranslator(null), predictorDevice)) {
+
+                // Pass in predictorDevice input
+                try (NDManager gpuManager = manager.newSubManager(predictorDevice)) {
+                    NDList input = new NDList(gpuManager.ones(new Shape(1, 10)));
+
+                    // The result should still be on predictorDevice
+                    NDList result = predictor.predict(input);
+                    Assert.assertEquals(result.get(0).getDevice(), predictorDevice);
+                }
+            }
+        }
+    }
+}
diff --git a/integration/src/main/java/ai/djl/integration/tests/inference/package-info.java b/integration/src/main/java/ai/djl/integration/tests/inference/package-info.java
new file mode 100644
index 000000000..d1e6ffa43
--- /dev/null
+++ b/integration/src/main/java/ai/djl/integration/tests/inference/package-info.java
@@ -0,0 +1,14 @@
+/*
+ * Copyright 2020 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 tests for {@link ai.djl.inference}. */
+package ai.djl.integration.tests.inference;