mirror of https://github.com/hpcaitech/ColossalAI
15 lines
487 B
Python
15 lines
487 B
Python
from .ophooks import BaseOpHook, register_ophooks_recursively
|
|
from .stateful_tensor import StatefulTensor
|
|
from .stateful_tensor_mgr import StatefulTensorMgr
|
|
from .tensor_placement_policy import AutoTensorPlacementPolicy, CPUTensorPlacementPolicy, CUDATensorPlacementPolicy
|
|
|
|
__all__ = [
|
|
"StatefulTensorMgr",
|
|
"StatefulTensor",
|
|
"CPUTensorPlacementPolicy",
|
|
"CUDATensorPlacementPolicy",
|
|
"AutoTensorPlacementPolicy",
|
|
"register_ophooks_recursively",
|
|
"BaseOpHook",
|
|
]
|