From 696c2aa9b7a80a1ec78529b0686988da3284da5d Mon Sep 17 00:00:00 2001
From: Miltos Allamanis <miallama@microsoft.com>
Date: Sat, 6 Jun 2020 08:26:05 +0100
Subject: [PATCH] Ensure optional dependencies are loaded only when necessary.

---
 README.md                                    | 9 +++++++--
 ptgnn/baseneuralmodel/abstractneuralmodel.py | 2 +-
 setup.py                                     | 2 +-
 3 files changed, 9 insertions(+), 4 deletions(-)

diff --git a/README.md b/README.md
index fb9f505..64722ef 100644
--- a/README.md
+++ b/README.md
@@ -56,7 +56,11 @@ CUDA 10.2).
 
 1. To install `ptgnn` from pypi, including all other dependencies:
     ```bash
-   pip install ptgpnn
+   pip install ptgnn
+    ```
+   If you want to use ptgnn sampels with Azure ML (e.g. the `--aml` flag in the implementation CLIs), install with
+   ```bash
+    pip install ptgnn[aml]
     ```
    or directly from the sources, `cd` into the root directory of the project and run
     ```bash
@@ -84,7 +88,8 @@ contact [opencode@microsoft.com](mailto:opencode@microsoft.com) with any additio
 
 
 ### Developing ptgnn
-![Run Tests](https://github.com/microsoft/ptgnn/workflows/Run%20Tests/badge.svg)
+![Unit Tests](https://github.com/microsoft/ptgnn/workflows/Unit%20Tests/badge.svg)
+[![Total alerts](https://img.shields.io/lgtm/alerts/g/microsoft/ptgnn.svg?logo=lgtm&logoWidth=18)](https://lgtm.com/projects/g/microsoft/ptgnn/alerts/)
 [![Code style: black](https://img.shields.io/badge/code%20style-black-000000.svg)](https://github.com/psf/black)
 
 To contribute to this library, first follow the next steps to setup your
diff --git a/ptgnn/baseneuralmodel/abstractneuralmodel.py b/ptgnn/baseneuralmodel/abstractneuralmodel.py
index b641a12..c946c59 100644
--- a/ptgnn/baseneuralmodel/abstractneuralmodel.py
+++ b/ptgnn/baseneuralmodel/abstractneuralmodel.py
@@ -20,7 +20,7 @@ from typing import (
 from typing_extensions import final
 
 import torch
-from dpu_utils.utils import BufferedIterator, ThreadedIterator, shuffled_iterator
+from dpu_utils.utils.iterators import BufferedIterator, ThreadedIterator, shuffled_iterator
 from torch import nn
 
 __all__ = ["AbstractNeuralModel"]
diff --git a/setup.py b/setup.py
index 3f1be1b..234e161 100644
--- a/setup.py
+++ b/setup.py
@@ -26,7 +26,7 @@ setup(
         "tqdm",
         "typing-extensions",
     ],
-    extras_require={"dev": ["black", "isort", "pre-commit"]},
+    extras_require={"dev": ["black", "isort", "pre-commit"], "aml": ["azureml"]},
     setup_requires=["setuptools_scm"],
     url="https://github.com/microsoft/ptgnn/",
     project_urls={