feat: display torch device on startup

This functionality disappeared at some point.
This commit is contained in:
psychedelicious
2024-03-27 17:28:06 +11:00
committed by Kent Keirsey
parent 4f2892d231
commit a291a42abc
2 changed files with 7 additions and 0 deletions

View File

@ -27,6 +27,9 @@ def choose_torch_device() -> torch.device:
else:
return torch.device(config.device)
def get_torch_device_name() -> str:
device = choose_torch_device()
return torch.cuda.get_device_name(device) if device.type == "cuda" else device.type.upper()
# We are in transition here from using a single global AppConfig to allowing multiple
# configurations. It is strongly recommended to pass the app_config to this function.