mirror of
https://github.com/vladmandic/sdnext.git
synced 2026-01-29 05:02:09 +03:00
9 lines
247 B
Python
9 lines
247 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)
|