vrr/heretic
2
0
Fork 0
mirror of https://github.com/p-e-w/heretic.git synced 2026-07-09 17:28:29 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

feat: enumerate all available GPUs on startup (#86)
Some checks failed
CI / Check and build (Python 3.10) (push) Has been cancelled
Details
CI / Check and build (Python 3.11) (push) Has been cancelled
Details
CI / Check and build (Python 3.12) (push) Has been cancelled
Details
CI / Check and build (Python 3.13) (push) Has been cancelled
Details

Browse source 
* feat: enumerate all available GPUs on startup

* feat: extend device enumeration to all accelerator types
This commit is contained in:
michaelh 2025-12-16 13:12:15 +01:00 • committed by GitHub
parent 243f821d93
commit 92d0c0d551
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 22 additions and 7 deletions
Download patch file Download diff file Expand all files Collapse all files

29
src/heretic/main.py
View file

@ -161,19 +161,34 @@ def run():
# Adapted from https://github.com/huggingface/accelerate/blob/main/src/accelerate/commands/env.py
if torch.cuda.is_available():
print(f"GPU type: [bold]{torch.cuda.get_device_name()}[/]")
count = torch.cuda.device_count()
print(f"Detected [bold]{count}[/] CUDA device(s):")
for i in range(count):
print(f"* GPU {i}: [bold]{torch.cuda.get_device_name(i)}[/]")
elif is_xpu_available():
print(f"XPU type: [bold]{torch.xpu.get_device_name()}[/]")
count = torch.xpu.device_count()
print(f"Detected [bold]{count}[/] XPU device(s):")
for i in range(count):
print(f"* XPU {i}: [bold]{torch.xpu.get_device_name(i)}[/]")
elif is_mlu_available():
print(f"MLU type: [bold]{torch.mlu.get_device_name()}[/]")
count = torch.mlu.device_count()
print(f"Detected [bold]{count}[/] MLU device(s):")
for i in range(count):
print(f"* MLU {i}: [bold]{torch.mlu.get_device_name(i)}[/]")
elif is_sdaa_available():
print(f"SDAA type: [bold]{torch.sdaa.get_device_name()}[/]")
count = torch.sdaa.device_count()
print(f"Detected [bold]{count}[/] SDAA device(s):")
for i in range(count):
print(f"* SDAA {i}: [bold]{torch.sdaa.get_device_name(i)}[/]")
elif is_musa_available():
print(f"MUSA type: [bold]{torch.musa.get_device_name()}[/]")
count = torch.musa.device_count()
print(f"Detected [bold]{count}[/] MUSA device(s):")
for i in range(count):
print(f"* MUSA {i}: [bold]{torch.musa.get_device_name(i)}[/]")
elif is_npu_available():
print(f"CANN version: [bold]{torch.version.cann}[/]")
print(f"NPU detected (CANN version: [bold]{torch.version.cann}[/])")
elif torch.backends.mps.is_available():
print("GPU type: [bold]Apple Metal (MPS)[/]")
print("Detected [bold]1[/] MPS device (Apple Metal)")
else:
print(
"[bold yellow]No GPU or other accelerator detected. Operations will be slow.[/]"