github/tinygrad
1
1
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-01-23 05:48:08 -05:00
Code Issues Packages Projects Releases Wiki Activity

nv default for ampere & ada (#5329)

Browse Source
This commit is contained in:
nimlgen
2024-07-08 19:01:27 +03:00
committed by GitHub
parent 51d6f372e4
commit bb2222e488
1 changed files with 1 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
tinygrad/device.py
View File

@@ -29,7 +29,7 @@ class _Device:
def DEFAULT(self) -> str:
device_from_env: Optional[str] = functools.reduce(lambda val, ele: ele if getenv(ele) == 1 else val, self._devices, None) # type: ignore
if device_from_env: return device_from_env
for device in ["METAL", "AMD", "CUDA", "GPU", "CLANG", "LLVM"]:
for device in ["METAL", "AMD", "NV", "CUDA", "GPU", "CLANG", "LLVM"]:
try:
if self[device]:
os.environ[device] = "1" # we set this in environment for spawned children