save/loadMetadata fixes

deepjavalibrary · Jul 1, 2021 · b780542 · b780542
1 parent 380fb3b
commit b780542
Showing 1 changed file with 4 additions and 12 deletions.
diff --git a/api/src/main/java/ai/djl/nn/norm/LayerNorm.java b/api/src/main/java/ai/djl/nn/norm/LayerNorm.java
@@ -165,26 +165,18 @@ public void prepare(Shape[] inputShapes) {
     @Override
     protected void saveMetadata(DataOutputStream os) throws IOException {
         saveInputShapes(os);
-        os.writeInt(normalizedShape.getShape().length);
-        for (int i = 0; i < normalizedShape.getShape().length; i++) {
-            os.writeLong(normalizedShape.getShape()[i]);
-        }
+        os.write(normalizedShape.getEncoded());
     }
 
     /** {@inheritDoc} */
     @Override
     public void loadMetadata(byte version, DataInputStream is)
             throws IOException, MalformedModelException {
-        if (version == VERSION) {
-            readInputShapes(is);
-        } else if (version != 1) {
+        if (version != VERSION) {
             throw new MalformedModelException("Unsupported encoding version: " + version);
         }
-        long[] shapeRaw = new long[is.readInt()];
-        for (int i = 0; i < shapeRaw.length; i++) {
-            shapeRaw[i] = is.readLong();
-        }
-        normalizedShape = new Shape(shapeRaw);
+        readInputShapes(is);
+        normalizedShape = Shape.decode(is);
     }
 
     /** The Builder to construct a {@link LayerNorm}. */