diff --git a/model/__init__.py b/model/__init__.py
index 7b09c41..202b11b 100644
--- a/model/__init__.py
+++ b/model/__init__.py
@@ -3,4 +3,4 @@ from model.dit import DiT
 from model.trainer import Trainer
 
 
-__all__ = ["CFM"]
+__all__ = ["CFM", "DiT", "Trainer"]