ColossalAI/tests/components_to_test/__init__.py

20 lines
440 B
Python

from . import (
beit,
bert,
gpt2,
hanging_param_model,
inline_op_model,
nested_model,
repeated_computed_layers,
resnet,
simple_net,
)
from .utils import run_fwd, run_fwd_bwd
from . import albert # isort:skip
__all__ = [
'bert', 'gpt2', 'hanging_param_model', 'inline_op_model', 'nested_model', 'repeated_computed_layers', 'resnet',
'simple_net', 'run_fwd_bwd', 'albert', 'beit', 'run_fwd'
]