fix: support single .h5 file loading in load_h5

2026-07-07 00:11:16 +08:00
parent 53ed52b4b8
commit 4af7acd449
1 changed files with 4 additions and 1 deletions
@@ -26,6 +26,9 @@ def load_h5(file_path: str, share_memory=True) -> Dict[str, List[Tensor]]:
    tensor_group: Dict[str, List[Tensor]] = {}

    root_path = Path(file_path)
+    if root_path.is_file() and root_path.suffix in (".h5", ".hdf5"):
+        h5_files = [root_path]
+    else:
        h5_files = list(root_path.rglob("*.h5")) + list(root_path.rglob("*.hdf5"))

    for h5_file in h5_files: