00c2c1f08b
Signed-off-by: Song Zhang <gepin.zs@antgroup.com> Co-authored-by: 22dimensions <waitingwind@foxmail.com>
18 lines
442 B
Python
18 lines
442 B
Python
import logging
|
|
from typing import Optional
|
|
|
|
import torch
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class DeviceConfig:
|
|
device: Optional[torch.device]
|
|
|
|
def __init__(self, device: str = "cuda") -> None:
|
|
if device in ["cuda", "xpu", "hpu", "cpu", "npu"]:
|
|
self.device_type = device
|
|
else:
|
|
raise RuntimeError(f"Not supported device type: {device}")
|
|
self.device = torch.device(self.device_type)
|