# from .model import DiT_Llama_5B_patch2
from .model import DiT_Llama_2B_patch2