mirror of
https://github.com/vladmandic/automatic
synced 2026-08-27 07:31:01 +02:00
Seunghoon Lee
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)
|