sirhcm/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests 2 Actions 8 Activity

cap default dataset wikipedia max_workers to 32 (#4345)

Browse Source 
64 on tinybox OOM
This commit is contained in:
chenyu
2024-04-28 21:55:21 -04:00
committed by GitHub
parent c1d8d425eb
commit 82d0ed3cf3
1 changed files with 1 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
extra/datasets/wikipedia.py
+1 -1
View File
@@ -402,7 +402,7 @@ if __name__ == "__main__":
elif sys.argv[1] == "pre-train":
os.makedirs(BASEDIR / "train", exist_ok=True)
if sys.argv[2] == "all": # Use all 500 parts for training generation
process_map(process_part, [part for part in range(500)], max_workers=getenv('NUM_WORKERS', os.cpu_count()), chunksize=1)
process_map(process_part, [part for part in range(500)], max_workers=getenv('NUM_WORKERS', min(os.cpu_count(), 32)), chunksize=1)
else: # Use a specific part for training generation
part = int(sys.argv[2])
os.makedirs(BASEDIR / "train" / str(part), exist_ok=True)