From 4cbd2cce9b41e888338a6f018ff873e85210f7ea Mon Sep 17 00:00:00 2001
From: liqunfu <liqfu@microsoft.com>
Date: Tue, 27 Apr 2021 09:20:21 -0700
Subject: [PATCH] . (#7466)

Co-authored-by: liqun <liqun@OrtTrainingDev4.af05slrtruoetgaxwwjv5nsq5e.px.internal.cloudapp.net>
---
 setup.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/setup.py b/setup.py
index 3126ca81dc..ca5ce95233 100644
--- a/setup.py
+++ b/setup.py
@@ -246,7 +246,8 @@ enable_training = parse_arg_remove_boolean(sys.argv, '--enable_training')
 if enable_training:
     packages.extend(['onnxruntime.training',
                      'onnxruntime.training.amp',
-                     'onnxruntime.training.optim'])
+                     'onnxruntime.training.optim',
+                     'onnxruntime.training.ortmodule'])
     requirements_file = "requirements-training.txt"
     # with training, we want to follow this naming convention:
     # stable: