mirror of
https://github.com/vladmandic/automatic
synced 2026-08-27 15:41:00 +02:00
Vladimir Mandic
10 lines
235 B
Python
10 lines
235 B
Python
from typing import Union
|
|
import torch
|
|
|
|
|
|
rDevice = Union[torch.device, int]
|
|
def get_device(device: rDevice | None=None) -> torch.device:
|
|
if device is None:
|
|
device = torch.dml.current_device()
|
|
return torch.device(device)
|