Merge pull request #279 from EleutherAI/ds-exists-check

Fix sweep ds config bug
EleutherAI · Jul 31, 2023 · 020f3fc · 020f3fc
2 parents ab7c504 + 4fc5c6a
commit 020f3fc
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/elk/training/sweep.py b/elk/training/sweep.py
@@ -22,7 +22,8 @@ def assert_models_exist(model_names):
 
 def assert_datasets_exist(dataset_names):
     for dataset_name in dataset_names:
-        get_dataset_config_info(dataset_name)
+        ds_name, _, config_name = dataset_name.partition(":")
+        get_dataset_config_info(ds_name, config_name=config_name)
 
 
 @dataclass