Skip to content

Commit b309304

Browse files
njzjzwanghan-iapcm
andauthored
fix: fix PyTorch model extension in simplify (#1596)
<!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit - **Bug Fixes** - Improved model file search logic to dynamically construct file patterns, ensuring better compatibility with different model suffixes. - Added an assertion to ensure at least one model file is found, providing clearer error messaging when no models are present. - **Tests** - Enhanced test setup by generating fake model files, improving robustness and reliability of tests. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: Jinzhe Zeng <[email protected]> Co-authored-by: Han Wang <[email protected]>
1 parent 3a1ed00 commit b309304

File tree

2 files changed

+6
-1
lines changed

2 files changed

+6
-1
lines changed

dpgen/simplify/simplify.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -221,7 +221,9 @@ def run_model_devi(iter_index, jdata, mdata):
221221
commands = []
222222
run_tasks = ["."]
223223
# get models
224-
models = glob.glob(os.path.join(work_path, "graph*pb"))
224+
suffix = _get_model_suffix(jdata)
225+
models = glob.glob(os.path.join(work_path, f"graph*{suffix}"))
226+
assert len(models) > 0, "No model file found."
225227
model_names = [os.path.basename(ii) for ii in models]
226228
task_model_list = []
227229
for ii in model_names:

tests/simplify/test_run_model_devi.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,9 @@ class TestOneH5(unittest.TestCase):
1717
def setUp(self):
1818
work_path = Path("iter.000000") / "01.model_devi"
1919
work_path.mkdir(parents=True, exist_ok=True)
20+
# fake models
21+
for ii in range(4):
22+
(work_path / f"graph.{ii:03d}.pb").touch()
2023
with tempfile.TemporaryDirectory() as tmpdir:
2124
with open(Path(tmpdir) / "test.xyz", "w") as f:
2225
f.write(

0 commit comments

Comments
 (0)