diff --git a/tinygrad/runtime/support/hcq.py b/tinygrad/runtime/support/hcq.py index 01cb6c69aa..66fb02cfda 100644 --- a/tinygrad/runtime/support/hcq.py +++ b/tinygrad/runtime/support/hcq.py @@ -62,9 +62,10 @@ if DEV.interface.startswith("MOCK"): from test.mockgpu.mockgpu import MockFileIO def hcq_filter_visible_devices(devs, device): assert (v:=getenv("HCQ_VISIBLE_DEVICES", "")) == "", f"HCQ_VISIBLE_DEVICES={v} is deprecated, use DEV={DEV.target(device, indices=v)} instead" - ids = [int(x) for x in DEV.target(device).indices.split(',') if x.strip()] + if '-' in (idstr:=DEV.target(device).indices): ids = list(range(int(idstr.split('-')[0]), int(idstr.split('-')[1])+1)) + else: ids = [int(x) for x in idstr.split(',') if x.strip()] assert all(x < len(devs) for x in ids), f"invalid visibility filter: {ids} ({pluralize('device', len(devs))} available)" - return [devs[x] for x in ids] if (ids:=[int(x) for x in DEV.target(device).indices.split(',') if x.strip()]) else devs + return [devs[x] for x in ids] if ids else devs SignalType = TypeVar('SignalType', bound='HCQSignal') HCQDeviceType = TypeVar('HCQDeviceType', bound='HCQCompiled')