mirror of https://github.com/hpcaitech/ColossalAI
10 lines
343 B
Python
10 lines
343 B
Python
from .nopadding_llama import NoPaddingLlamaModelInferPolicy
|
|
from .padding_llama import PaddingLlamaModelInferPolicy
|
|
|
|
model_policy_map = {
|
|
"padding_llama": PaddingLlamaModelInferPolicy,
|
|
"nopadding_llama": NoPaddingLlamaModelInferPolicy,
|
|
}
|
|
|
|
__all__ = ["PaddingLlamaModelInferPolicy", "NoPaddingLlamaModelInferPolicy", "model_polic_map"]
|