allow empty batch data (#1569)

* allow empty batch data * Fix java code format Change-Id: I2d1e2894d37fb4f6a09ad3c70344bd1e822da65c Co-authored-by: Frank Liu <frankfliu2000@gmail.com>
deepjavalibrary · Apr 14, 2022 · ef7d6ca · ef7d6ca
1 parent 94d2765
commit ef7d6ca
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/api/src/main/java/ai/djl/nn/Blocks.java b/api/src/main/java/ai/djl/nn/Blocks.java
@@ -29,6 +29,10 @@ private Blocks() {}
      */
     public static NDArray batchFlatten(NDArray array) {
         long batch = array.size(0);
+        if (batch == 0) {
+            // calculate the size of second dimension manually as using -1 would not work here
+            return array.reshape(batch, array.getShape().slice(1).size());
+        }
         return array.reshape(batch, -1);
     }