fix loaded config in test

RaymondLi0 · RaymondLi0 · commit 229050d0a91a · 2024-12-06T19:11:58.000Z
diff --git a/tests/test_checkpoint.py b/tests/test_checkpoint.py
@@ -241,7 +241,7 @@ def _compare_configs(config_ref, config_test):
 @pytest.mark.depends(on=["test_converted_distributed"])
 def test_load_pretrained_distributed_checkpoint():
     config = TEST_ARCHITECTURE_CONFIG_CLS.from_dict(
-        yaml.safe_load((_CKPT_PATH / ".." / ".." / "config.yaml").open("r")), strict=False
+        yaml.safe_load((_CKPT_PATH / ".." / ".." / "config.yaml").open("r"))["model"]["base_model"], strict=False
     )
     pretrained_config_ref = CheckpointLoadConfig(
         path=_CKPT_PATH,

Original file line number	Diff line number	Diff line change
`@@ -241,7 +241,7 @@ def _compare_configs(config_ref, config_test):`
`241`	`241`	`@pytest.mark.depends(on=["test_converted_distributed"])`
`242`	`242`	`def test_load_pretrained_distributed_checkpoint():`
`243`	`243`	`config = TEST_ARCHITECTURE_CONFIG_CLS.from_dict(`
`244`		`- yaml.safe_load((_CKPT_PATH / ".." / ".." / "config.yaml").open("r")), strict=False`
	`244`	`+ yaml.safe_load((_CKPT_PATH / ".." / ".." / "config.yaml").open("r"))["model"]["base_model"], strict=False`
`245`	`245`	`)`
`246`	`246`	`pretrained_config_ref = CheckpointLoadConfig(`
`247`	`247`	`path=_CKPT_PATH,`