fix some bugs in pytorch based examples (#2009)

deepjavalibrary · Sep 13, 2022 · 5073966 · 5073966
1 parent 93ec732
commit 5073966
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 5 deletions.
diff --git a/api/src/main/java/ai/djl/modality/cv/translator/BigGANTranslator.java b/api/src/main/java/ai/djl/modality/cv/translator/BigGANTranslator.java
@@ -40,15 +40,13 @@ public BigGANTranslator(float truncation) {
     /** {@inheritDoc} */
     @Override
     public Image[] processOutput(TranslatorContext ctx, NDList list) {
-        NDArray output = list.get(0).addi(1).muli(128).clip(0, 255);
-
+        NDArray output = list.get(0).duplicate().addi(1).muli(128).clip(0, 255);
         int sampleSize = (int) output.getShape().get(0);
         Image[] images = new Image[sampleSize];
 
         for (int i = 0; i < sampleSize; ++i) {
             images[i] = ImageFactory.getInstance().fromNDArray(output.get(i));
         }
-
         return images;
     }
 

diff --git a/api/src/main/java/ai/djl/modality/cv/translator/StyleTransferTranslator.java b/api/src/main/java/ai/djl/modality/cv/translator/StyleTransferTranslator.java
@@ -34,7 +34,7 @@ public NDList processInput(TranslatorContext ctx, Image input) {
     /** {@inheritDoc} */
     @Override
     public Image processOutput(TranslatorContext ctx, NDList list) {
-        NDArray output = list.get(0).addi(1).muli(128);
+        NDArray output = list.get(0).duplicate().addi(1).muli(128);
         return ImageFactory.getInstance().fromNDArray(output.squeeze());
     }
 

diff --git a/examples/src/main/java/ai/djl/examples/inference/StyleTransfer.java b/examples/src/main/java/ai/djl/examples/inference/StyleTransfer.java
@@ -52,7 +52,9 @@ public static void main(String[] args) throws IOException, ModelException, Trans
         Image input = ImageFactory.getInstance().fromFile(Paths.get(imagePath));
         Image output = transfer(input, artist);
 
-        logger.info("Using PyTorch Engine. " + artist + " painting generated.");
+        logger.info(
+                "Using PyTorch Engine. {} painting generated. Image saved in build/output/cyclegan",
+                artist);
         save(output, artist.toString(), "build/output/cyclegan/");
     }