ColossalAI/colossalai/accelerator/__init__.py

14 lines
365 B
Python

from .api import auto_set_accelerator, get_accelerator, set_accelerator
from .base_accelerator import BaseAccelerator
from .cuda_accelerator import CudaAccelerator
from .npu_accelerator import NpuAccelerator
__all__ = [
"get_accelerator",
"set_accelerator",
"auto_set_accelerator",
"BaseAccelerator",
"CudaAccelerator",
"NpuAccelerator",
]