From 2558586fb01e07c48f5edd8df979651b4cdb7fb6 Mon Sep 17 00:00:00 2001
From: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>
Date: Wed, 7 Jun 2023 20:34:53 -0400
Subject: [PATCH] fix customized data shape in deepmd/hdf5

`deepmd/npy` and `deepmd/raw` are ok, and only `deepmd/hdf5` has this bug.
---
 dpdata/deepmd/hdf5.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/dpdata/deepmd/hdf5.py b/dpdata/deepmd/hdf5.py
index b4ae1a3c6..f2ee9c8a2 100644
--- a/dpdata/deepmd/hdf5.py
+++ b/dpdata/deepmd/hdf5.py
@@ -118,11 +118,14 @@ def to_system_data(
                 f"Shape of {dtype.name} is not (nframes, ...), but {dtype.shape}. This type of data will not converted from deepmd/hdf5 format."
             )
             continue
+        shape = [
+            natoms if xx == dpdata.system.Axis.NATOMS else xx for xx in dtype.shape[1:]
+        ]
 
         data_types[dtype.name] = {
             "fn": dtype.name,
             "labeled": True,
-            "shape": dtype.shape[1:],
+            "shape": shape,
             "required": False,
         }