fix: 修复 pipeline 模块中的打包逻辑缺陷并完善测试覆盖

pipeline/io.py

@@ -28,9 +28,9 @@ class IOHandler:
         return folders
 
     @staticmethod
-    def save_h5(file_path: str, file_name: str, tensor_group: Dict[str, List[Tensor]]) -> None:
-        os.makedirs(file_path, exist_ok=True)
-        full_path = os.path.join(file_path, f"{file_name}.h5")
+    def save_h5(output_dir: str, file_name: str, tensor_group: Dict[str, List[Tensor]]) -> None:
+        os.makedirs(output_dir, exist_ok=True)
+        full_path = os.path.join(output_dir, f"{file_name}.h5")
         with h5py.File(full_path, 'w') as f:
             for key, tensors in tensor_group.items():
                 grp = f.create_group(key)
@@ -38,19 +38,26 @@ class IOHandler:
                     grp.create_dataset(f'data_{idx}', data=tensor.cpu().numpy())
 
     @staticmethod
-    def load_h5(file_path: str, share_memory: bool = True) -> Dict[str, List[Tensor]]:
+    def load_h5(file_path: str, share_memory=True) -> Dict[str, List[Tensor]]:
         tensor_group: Dict[str, List[Tensor]] = {}
+
         root_path = Path(file_path)
         h5_files = list(root_path.rglob("*.h5")) + list(root_path.rglob("*.hdf5"))
+        
         for h5_file in h5_files:
             with h5py.File(h5_file, 'r') as f:
                 for key in f.keys():
                     grp = f[key]
-                    tensors = [
-                        (torch.from_numpy(dset[:]).share_memory_() if share_memory
-                         else torch.from_numpy(dset[:]))
-                        for dset_name in grp.keys()
-                        for dset in [grp[dset_name]]
-                    ]
-                    tensor_group.setdefault(key, []).extend(tensors)
-        return tensor_group
+                    dsets = []
+                    for dset_name in grp.keys():
+                        dset = grp[dset_name]
+                        tensor = torch.from_numpy(dset[:])
+                        if share_memory:
+                            tensor = tensor.share_memory_()
+                        dsets.append(tensor)
+                
+                    if tensor_group.get(key) is None:
+                        tensor_group[key] = []
+                    tensor_group[key].extend(dsets)
+
+        return tensor_group