ViperEkura/DataPipeline
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

refactor: 更新dump_files函数

Browse Source
This commit is contained in:
ViperEkura
2026-02-27 19:46:00 +08:00
parent eb2c6d4ffb
commit 3a8edc156d
1 changed files with 5 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
modules/utils.py
+5 -3
View File
@@ -129,13 +129,15 @@ def dump_files(
list_tensor = [arrow[key] for arrow in arrows] list_tensor = [arrow[key] for arrow in arrows]
package[key] = list_tensor package[key] = list_tensor
output_package: Dict[str, Tensor] = {} output_package: Dict[str, List[Tensor]] = {}
for key in output_keys: for key in output_keys:
if packing_size > 0: if packing_size > 0:
print(f"Packaging key: '{key}'") print(f"Packaging key: '{key}'")
package[key] = pack_sequences(package[key], packing_size, pad_value) sequence = pack_sequences(package[key], packing_size, pad_value)
sequence = torch.cat(package[key]) else:
sequence = package[key]
output_package[key] = sequence output_package[key] = sequence
save_h5(out_file_path, output_package) save_h5(out_file_path, output_package)