diff --git a/_doc/bench/bench_orttraining_nn_gpu.py b/_doc/bench/bench_orttraining_nn_gpu.py
index 8d70a6c4..ad8842de 100644
--- a/_doc/bench/bench_orttraining_nn_gpu.py
+++ b/_doc/bench/bench_orttraining_nn_gpu.py
@@ -35,8 +35,9 @@
 from sklearn.neural_network import MLPRegressor
 from sklearn.metrics import mean_squared_error
 from mlprodict.onnx_conv import to_onnx
-from onnxcustom.training import add_loss_output, get_train_initializer
-from onnxcustom.training.optimizers import OrtGradientOptimizer
+from onnxcustom.training import (
+    add_loss_output, get_train_initializer,
+    OrtGradientOptimizer)
 
 
 def benchmark(N=1000, n_features=20, hidden_layer_sizes="25,25", max_iter=1000,
diff --git a/_doc/examples/plot_orttraining_benchmark.py b/_doc/examples/plot_orttraining_benchmark.py
index 867b90df..cf6eb795 100644
--- a/_doc/examples/plot_orttraining_benchmark.py
+++ b/_doc/examples/plot_orttraining_benchmark.py
@@ -28,8 +28,9 @@
 from sklearn.model_selection import train_test_split
 from sklearn.neural_network import MLPRegressor
 from mlprodict.onnx_conv import to_onnx
-from onnxcustom.training import add_loss_output, get_train_initializer
-from onnxcustom.training.optimizers import OrtGradientOptimizer
+from onnxcustom.training import (
+    add_loss_output, get_train_initializer,
+    OrtGradientOptimizer)
 
 
 X, y = make_regression(2000, n_features=100, bias=2)
diff --git a/_unittests/ut_training/test_data_loader.py b/_unittests/ut_training/test_data_loader.py
index d5ed2e80..5cefebb5 100644
--- a/_unittests/ut_training/test_data_loader.py
+++ b/_unittests/ut_training/test_data_loader.py
@@ -7,7 +7,7 @@
 from pyquickhelper.pycode import ExtTestCase
 from sklearn.datasets import make_regression
 from onnxruntime import OrtValue
-from onnxcustom.training.data_loader import OrtDataLoader
+from onnxcustom.training import OrtDataLoader
 
 
 class TestDataLoadeer(ExtTestCase):
diff --git a/_unittests/ut_training/test_optimizers.py b/_unittests/ut_training/test_optimizers.py
index 5e81f2d8..cd0ad066 100644
--- a/_unittests/ut_training/test_optimizers.py
+++ b/_unittests/ut_training/test_optimizers.py
@@ -22,8 +22,8 @@ class TestOptimizers(ExtTestCase):
 
     @unittest.skipIf(TrainingSession is None, reason="not training")
     def test_ort_gradient_optimizers_use_numpy(self):
-        from onnxcustom.training.orttraining import add_loss_output
-        from onnxcustom.training.optimizers import OrtGradientOptimizer
+        from onnxcustom.training import (
+            add_loss_output, OrtGradientOptimizer)
         X, y = make_regression(  # pylint: disable=W0632
             100, n_features=10, bias=2)
         X = X.astype(numpy.float32)
@@ -78,8 +78,8 @@ def test_ort_gradient_optimizers_use_ort(self):
 
     @unittest.skipIf(TrainingSession is None, reason="not training")
     def test_ort_gradient_optimizers_optimal_use_numpy(self):
-        from onnxcustom.training.orttraining import add_loss_output
-        from onnxcustom.training.optimizers import OrtGradientOptimizer
+        from onnxcustom.training import (
+            add_loss_output, OrtGradientOptimizer)
         X, y = make_regression(  # pylint: disable=W0632
             100, n_features=10, bias=2)
         X = X.astype(numpy.float32)
@@ -134,8 +134,8 @@ def test_ort_gradient_optimizers_optimal_use_ort(self):
 
     @unittest.skipIf(TrainingSession is None, reason="not training")
     def test_ort_gradient_optimizers_evaluation_use_numpy(self):
-        from onnxcustom.training.orttraining import add_loss_output
-        from onnxcustom.training.optimizers import OrtGradientOptimizer
+        from onnxcustom.training import (
+            add_loss_output, OrtGradientOptimizer)
         X, y = make_regression(  # pylint: disable=W0632
             100, n_features=10, bias=2)
         X = X.astype(numpy.float32)
@@ -164,8 +164,8 @@ def test_ort_gradient_optimizers_evaluation_use_numpy(self):
 
     @unittest.skipIf(TrainingSession is None, reason="not training")
     def test_ort_gradient_optimizers_evaluation_use_ort(self):
-        from onnxcustom.training.orttraining import add_loss_output
-        from onnxcustom.training.optimizers import OrtGradientOptimizer
+        from onnxcustom.training import (
+            add_loss_output, OrtGradientOptimizer)
         X, y = make_regression(  # pylint: disable=W0632
             100, n_features=10, bias=2)
         X = X.astype(numpy.float32)
diff --git a/onnxcustom/training/__init__.py b/onnxcustom/training/__init__.py
index 22f037b2..bda5a961 100644
--- a/onnxcustom/training/__init__.py
+++ b/onnxcustom/training/__init__.py
@@ -1,6 +1,8 @@
 """
 @file
-@brief Shortcuts to *training*.
+@brief Shortcuts to *orttraining*.
 """
 
+from .data_loader import OrtDataLoader  # noqa
+from .optimizers import OrtGradientOptimizer  # noqa
 from .orttraining import add_loss_output, get_train_initializer  # noqa
diff --git a/onnxcustom/training/optimizers.py b/onnxcustom/training/optimizers.py
index 23789179..c5c8b3e6 100644
--- a/onnxcustom/training/optimizers.py
+++ b/onnxcustom/training/optimizers.py
@@ -1,6 +1,7 @@
 """
 @file
-@brief Helper for :epkg:`onnxruntime-training`.
+@brief Train a machine learned model
+with :epkg:`onnxruntime-training`.
 """
 import inspect
 import numpy
diff --git a/onnxcustom/training/ortgradient.py b/onnxcustom/training/ortgradient.py
new file mode 100644
index 00000000..98be01bc
--- /dev/null
+++ b/onnxcustom/training/ortgradient.py
@@ -0,0 +1,4 @@
+"""
+@file
+@brief Helpers for :epkg:`onnxruntime-training`.
+"""
diff --git a/onnxcustom/training/orttraining.py b/onnxcustom/training/orttraining.py
index 86e31f34..8a5e1515 100644
--- a/onnxcustom/training/orttraining.py
+++ b/onnxcustom/training/orttraining.py
@@ -1,6 +1,6 @@
 """
 @file
-@brief Helper for :epkg:`onnxruntime-training`.
+@brief Manipulate ONNX graph to train a model.
 """
 from onnx.helper import (
     make_node, make_graph, make_model, make_tensor_value_info,