dump the deepmd/npy data with only one set

deepmodeling · Feb 23, 2021 · dc17bfc · dc17bfc
1 parent f995ee5
commit dc17bfc
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/dpgen/collect/collect.py b/dpgen/collect/collect.py
@@ -78,7 +78,9 @@ def collect_data(target_folder, param_file, output,
     # dump iter data
     for kk in coll_data.keys():
         out_dir = 'sys.%s' % kk
-        coll_data[kk].to('deepmd/npy', os.path.join(output, out_dir))
+        nframes = coll_data[kk].get_nframes()
+        coll_data[kk].to('deepmd/npy', os.path.join(output, out_dir), set_size = nframes)
+        # coll_data[kk].to('deepmd/npy', os.path.join(output, out_dir))
 
 def gen_collect(args):
     collect_data(args.JOB_DIR, args.parameter, args.OUTPUT,