add dataset name to download.py

2021-11-23 07:48:43 -08:00
parent 7e68d4f712
commit 291ae1a1ee
2 changed files with 7 additions and 2 deletions
--- a/download.py
+++ b/download.py
@@ -4,7 +4,12 @@ import sys

 import datasets

-mats = datasets.ALL_REAL_MATS
+dataset = None
+for ds in datasets.DATASETS:
+    if ds.name == sys.argv[1]:
+        dataset = ds
+        break
+mats = dataset.mats

 print(len(mats))