mirror of
https://github.com/vladmandic/sdnext.git
synced 2026-01-27 15:02:48 +03:00
10 lines
248 B
Python
10 lines
248 B
Python
from typing import Optional, Union
|
|
import torch
|
|
|
|
|
|
rDevice = Union[torch.device, int]
|
|
def get_device(device: Optional[rDevice]=None) -> torch.device:
|
|
if device is None:
|
|
device = torch.dml.current_device()
|
|
return torch.device(device)
|