fixing test

argonne-lcf · Feb 6, 2025 · 0c86f71 · 0c86f71
1 parent 58c220e
commit 0c86f71
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tests/dlio_benchmark_test.py b/tests/dlio_benchmark_test.py
@@ -262,7 +262,7 @@ def test_checkpoint_epoch(framework, model_size, optimizers, num_layers, layer_p
         comm.Barrier()
         benchmark = run_benchmark(cfg)
         output = pathlib.Path("./checkpoints")
-        load_bin = list(output.glob("*"))
+        load_bin = list(output.glob("*/*.pt"))
         n = 0
         if len(layer_params) > 0:
             n = num_layers