fix bug in dplr_v2_to_v3

chenggroup · Dec 10, 2024 · d505495 · d505495
1 parent db8f514
commit d505495
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/ai2_kit/domain/dplr.py b/ai2_kit/domain/dplr.py
@@ -341,8 +341,9 @@ def dplr_v2_to_v3(data_path: str, sel_symbol: list):
             sel_ids = np.where(np.isin(symbols, sel_symbol))[0]
             n_atoms = len(atype)
 
-            raw_data = np.load(fname).reshape([n_frames, len(sel_ids), -1])
+            raw_data = np.load(fname)
             n_frames = raw_data.shape[0]
+            raw_data = np.reshape(raw_data, [n_frames, len(sel_ids), -1])
             n_dim = raw_data.shape[2]
 
             full_data = np.zeros([n_frames, n_atoms, n_dim])