diff --git a/python/kubeflow/trainer/__init__.py b/python/kubeflow/trainer/__init__.py
index e8f452d0..827f4933 100644
--- a/python/kubeflow/trainer/__init__.py
+++ b/python/kubeflow/trainer/__init__.py
@@ -40,7 +40,8 @@
 )
 
 __all__ = [
-    "BuiltinTrainer", "CustomTrainer", "DataFormat", "DataType", "Framework",
-    "HuggingFaceDatasetInitializer", "HuggingFaceModelInitializer", "Initializer",
-    "Loss", "Runtime", "TorchTuneConfig", "TorchTuneInstructDataset", "Trainer", "TrainerType"
+    "BuiltinTrainer", "CustomTrainer", "DataFormat", "DATASET_PATH", "DataType", "Framework",
+    "HuggingFaceDatasetInitializer", "HuggingFaceModelInitializer", "Initializer", "Loss",
+    "MODEL_PATH", "Runtime", "TorchTuneConfig", "TorchTuneInstructDataset", "Trainer",
+    "TrainerClient", "TrainerType"
 ]