diff --git a/diabetes_regression/training/train_aml.py b/diabetes_regression/training/train_aml.py
index 5bf76cb4..18cfb852 100644
--- a/diabetes_regression/training/train_aml.py
+++ b/diabetes_regression/training/train_aml.py
@@ -23,7 +23,8 @@
 ARISING IN ANY WAY OUT OF THE USE OF THE SOFTWARE CODE, EVEN IF ADVISED OF THE
 POSSIBILITY OF SUCH DAMAGE.
 """
-from azureml.core.run import Dataset, Datastore, Run, Workspace
+from azureml.core.run import Run
+from azureml.core import Dataset, Datastore, Workspace
 import os
 import argparse
 import joblib