diff --git a/tinygrad/engine/search.py b/tinygrad/engine/search.py index 368f9963d9..c2906a29bf 100644 --- a/tinygrad/engine/search.py +++ b/tinygrad/engine/search.py @@ -142,7 +142,7 @@ def beam_search(lin:Kernel, rawbufs:list[Buffer], amt:int, allow_test_size=True, beam: list[tuple[Kernel, float]] = [(lin, float("inf"))] seen_libs = set() - default_parallel = multiprocessing.cpu_count() if lin.opts.device in {"CUDA", "AMD", "NV", "METAL"} else 0 + default_parallel = multiprocessing.cpu_count() if lin.opts.device in {"CUDA", "AMD", "NV", "METAL", "HIP"} else 0 if beam_pool is None and (workers := getenv("PARALLEL", default_parallel)): beam_pool = multiprocessing.get_context("spawn").Pool(workers, _init_worker, (), getenv("BEAM_MAX_TASKS_PER_CHILD", 16)) @atexit.register