15 lines
328 B
Python
15 lines
328 B
Python
|
from dataclasses import dataclass
|
||
|
|
|
||
|
|
import torch_npu
|
||
|
|
|
||
|
|
|
||
|
|
@dataclass
|
||
|
|
class NPUSocInfo:
|
||
|
|
is_a3: bool = False
|
||
|
|
|
||
|
|
def __post_init__(self):
|
||
|
|
torch_npu.npu._lazy_init()
|
||
|
|
self.soc_version = torch_npu._C._npu_get_soc_version()
|
||
|
|
if self.soc_version in (250, 251, 252, 253, 254, 255):
|
||
|
|
self.is_a3 = True
|