Skip to content

Commit

Permalink
allow empty batch data (#1569)
Browse files Browse the repository at this point in the history
* allow empty batch data

* Fix java code format

Change-Id: I2d1e2894d37fb4f6a09ad3c70344bd1e822da65c

Co-authored-by: Frank Liu <frankfliu2000@gmail.com>
  • Loading branch information
patins1 and frankfliu authored Apr 14, 2022
1 parent 94d2765 commit ef7d6ca
Showing 1 changed file with 4 additions and 0 deletions.
4 changes: 4 additions & 0 deletions api/src/main/java/ai/djl/nn/Blocks.java
Original file line number Diff line number Diff line change
Expand Up @@ -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);
}

Expand Down

0 comments on commit ef7d6ca

Please sign in to comment.