Repository: pygod-team/pygod (language: python)

Issue:
Add conad
Add the implementation of conad in Contrastive Attributed Network Anomaly Detection with Data Augmentation (PAKDD 22).
testing for Guide is slow
any thoughts for speeding this up?


Discussion / hints:

not critical but can be accelerated maybe
It's because of the high time complexity of VF2 algorithm counting node motif degree in preprocessing. Potential improvement is provided in #13.

Fix patch:
diff --git a/CHANGES.txt b/CHANGES.txt
index 3f9245f..629681f 100644
--- a/CHANGES.txt
+++ b/CHANGES.txt
@@ -1,2 +1,7 @@
 v<0.1.0>, <04/01/2022> -- Initial release.
 v<0.1.1>, <04/04/2022> -- Improvements on stability.
+v<0.2.0>, <04/14/2022> -- Support mini-batch and sampling.
+v<0.2.0>, <04/23/2022> -- Unify the default device choice with robust check.
+v<0.2.0>, <04/23/2022> -- Add CoLA, ANEMONE, CONAD.
+
+
diff --git a/README.rst b/README.rst
index 5233d01..0b415b1 100644
--- a/README.rst
+++ b/README.rst
@@ -40,19 +40,17 @@ suspicious activities in social networks [#Dou2020Enhancing]_  and security syst
 
 PyGOD includes more than **10** latest graph-based detection algorithms,
 such as DOMINANT (SDM'19) and GUIDE (BigData'21).
-For consistently and accessibility, PyGOD is developed on top of `PyTorch Geometric (PyG) <https://www.pyg.org/>`_
+For consistency and accessibility, PyGOD is developed on top of `PyTorch Geometric (PyG) <https://www.pyg.org/>`_
 and `PyTorch <https://pytorch.org/>`_, and follows the API design of `PyOD <https://github.com/yzhao062/pyod>`_.
 See examples below for detecting outliers with PyGOD in 5 lines!
 
-**PyGOD** is under actively developed and will be updated frequently!
-Please **star**, **watch**, and **fork**.
-
 
 **PyGOD is featured for**:
 
 * **Unified APIs, detailed documentation, and interactive examples** across various graph-based algorithms.
 * **Comprehensive coverage** of more than 10 latest graph outlier detectors.
-* **Full support of detections at multiple levels**, such as node-, edge-, and graph-level tasks (WIP).
+* **Full support of detections at multiple levels**, such as node-, edge- (WIP), and graph-level tasks (WIP).
+* **Scalable design for processing large graphs** via mini-batch and sampling.
 * **Streamline data processing with PyG**--fully compatible with PyG data objects.
 
 **Outlier Detection Using PyGOD with 5 Lines of Code**\ :
@@ -63,31 +61,33 @@ Please **star**, **watch**, and **fork**.
     # train a dominant detector
     from pygod.models import DOMINANT
 
-    model = DOMINANT()  # hyperparameters can be set here
+    model = DOMINANT(num_layers=4, epoch=20)  # hyperparameters can be set here
     model.fit(data)  # data is a Pytorch Geometric data object
 
     # get outlier scores on the input data
     outlier_scores = model.decision_scores # raw outlier scores on the input data
 
-    # predict on the new data
+    # predict on the new data in the inductive setting
     outlier_scores = model.decision_function(test_data) # raw outlier scores on the input data  # predict raw outlier scores on test
 
-**Citing PyGOD (to be announced soon)**\ :
 
-`PyGOD paper <https://pygod.org>`_ will be available on arxiv soon.
+**Citing PyGOD**\ :
+
+`PyGOD paper <https://arxiv.org/abs/2204.12095>`_ is available on arxiv.
 If you use PyGOD in a scientific publication, we would appreciate
-citations to the following paper (to be announced)::
+citations to the following paper::
 
-    @article{tba,
-      author  = {tba},
-      title   = {PyGOD: A Comprehensive Python Library for Graph Outlier Detection},
-      journal = {tba},
+    @article{pygod2022,
+      author  = {Liu, Kay and Dou, Yingtong and Zhao, Yue and Ding, Xueying and Hu, Xiyang and Zhang, Ruitong and Ding, Kaize and Chen, Canyu and Peng, Hao and Shu, Kai and Chen, George H. and Jia, Zhihao and Yu, Philip S.},
+      title   = {PyGOD: A Python Library for Graph Outlier Detection},
+      journal = {arXiv preprint arXiv:2204.12095},
       year    = {2022},
     }
 
 or::
 
-    tba, 2022. PyGOD: A Comprehensive Python Library for Graph Outlier Detection. tba.
+    Liu, K., Dou, Y., Zhao, Y., Ding, X., Hu, X., Zhang, R., Ding, K., Chen, C., Peng, H., Shu, K., Chen, G.H., Jia, Z., and Yu, P.S. 2022. PyGOD: A Python Library for Graph Outlier Detection. arXiv preprint arXiv:2204.12095.
+
 
 
 ----
@@ -114,11 +114,9 @@ Alternatively, you could clone and run setup.py file:
 **Required Dependencies**\ :
 
 * Python 3.6 +
-* argparse>=1.4.0
 * numpy>=1.19.4
 * scikit-learn>=0.22.1
 * scipy>=1.5.2
-* pandas>=1.1.3
 * setuptools>=50.3.1.post20201107
 
 
@@ -139,21 +137,25 @@ API Cheatsheet & Reference
 
 Full API Reference: (https://docs.pygod.org). API cheatsheet for all detectors:
 
-
 * **fit(X)**\ : Fit detector.
 * **decision_function(G)**\ : Predict raw anomaly score of PyG data G using the fitted detector.
-* **predict(G)**\ : Predict if nodes in PyG data G is an outlier or not using the fitted detector.
-* **predict_proba(G)**\ : Predict the probability of nodes in PyG data G being outlier using the fitted detector.
-* **predict_confidence(G)**\ : Predict the model's node-wise confidence (available in predict and predict_proba) [#Perini2020Quantifying]_.
-
 
 Key Attributes of a fitted model:
 
-
 * **decision_scores_**\ : The outlier scores of the training data. The higher, the more abnormal.
   Outliers tend to have higher scores.
 * **labels_**\ : The binary labels of the training data. 0 stands for inliers and 1 for outliers/anomalies.
 
+For the inductive setting:
+
+* **predict(G)**\ : Predict if nodes in PyG data G is an outlier or not using the fitted detector.
+* **predict_proba(G)**\ : Predict the probability of nodes in PyG data G being outlier using the fitted detector.
+* **predict_confidence(G)**\ : Predict the model's node-wise confidence (available in predict and predict_proba) [#Perini2020Quantifying]_.
+
+
+**Input of PyGOD**: Please pass in a `PyTorch Geometric (PyG) <https://www.pyg.org/>`_ data object.
+See `PyG data processing examples <https://pytorch-geometric.readthedocs.io/en/latest/notes/introduction.html#data-handling-of-graphs>`_.
+
 
 Implemented Algorithms
 ^^^^^^^^^^^^^^^^^^^^^^
@@ -162,20 +164,23 @@ PyGOD toolkit consists of two major functional groups:
 
 **(i) Node-level detection** :
 
-===================  ===================  ==================  ======================================================================================================  =====  ========================================
-Type                 Backbone             Abbr                Algorithm                                                                                               Year   Ref
-===================  ===================  ==================  ======================================================================================================  =====  ========================================
-Unsupervised         MLP                  MLPAE               Anomaly Detection Using Autoencoders with Nonlinear Dimensionality Reduction                            2014   [#Sakurada2014Anomaly]_
-Unsupervised         GNN                  GCNAE               Variational Graph Auto-Encoders                                                                         2016   [#Kipf2016Variational]_
-Unsupervised         MF                   ONE                 Outlier aware network embedding for attributed networks                                                 2019   [#Bandyopadhyay2019Outlier]_
-Unsupervised         GNN                  DOMINANT            Deep anomaly detection on attributed networks                                                           2019   [#Ding2019Deep]_
-Unsupervised         GNN                  DONE                Outlier Resistant Unsupervised Deep Architectures for Attributed Network Embedding                      2020   [#Bandyopadhyay2020Outlier]_
-Unsupervised         GNN                  AdONE               Outlier Resistant Unsupervised Deep Architectures for Attributed Network Embedding                      2020   [#Bandyopadhyay2020Outlier]_
-Unsupervised         GNN                  AnomalyDAE          AnomalyDAE: Dual autoencoder for anomaly detection on attributed networks                               2020   [#Fan2020AnomalyDAE]_
-Unsupervised         GAN                  GAAN                Generative Adversarial Attributed Network Anomaly Detection                                             2020   [#Chen2020Generative]_
-Unsupervised         GNN                  OCGNN               One-Class Graph Neural Networks for Anomaly Detection in Attributed Networks                            2021   [#Wang2021One]_
-Unsupervised         GNN                  GUIDE               Higher-order Structure Based Anomaly Detection on Attributed Networks                                   2021   [#Yuan2021Higher]_
-===================  ===================  ==================  ======================================================================================================  =====  ========================================
+===================  ===================  ==================  =====  ===========  ========================================
+Type                 Backbone             Abbr                Year   Sampling      Ref
+===================  ===================  ==================  =====  ===========  ========================================
+Unsupervised         MLP                  MLPAE               2014   Yes          [#Sakurada2014Anomaly]_
+Unsupervised         GNN                  GCNAE               2016   Yes          [#Kipf2016Variational]_
+Unsupervised         MF                   ONE                 2019   No           [#Bandyopadhyay2019Outlier]_
+Unsupervised         GNN                  DOMINANT            2019   Yes          [#Ding2019Deep]_
+Unsupervised         GNN                  DONE                2020   Yes          [#Bandyopadhyay2020Outlier]_
+Unsupervised         GNN                  AdONE               2020   Yes          [#Bandyopadhyay2020Outlier]_
+Unsupervised         GNN                  AnomalyDAE          2020   Yes          [#Fan2020AnomalyDAE]_
+Unsupervised         GAN                  GAAN                2020   Yes          [#Chen2020Generative]_
+Unsupervised         GNN                  OCGNN               2021   Yes          [#Wang2021One]_
+Unsupervised/SSL     GNN                  CoLA (beta)         2021   In progress  [#Liu2021Anomaly]_
+Unsupervised/SSL     GNN                  ANEMONE (beta)      2021   In progress  [#Jin2021ANEMONE]_
+Unsupervised         GNN                  GUIDE               2021   Yes          [#Yuan2021Higher]_
+Unsupervised/SSL     GNN                  CONAD               2022   Yes          [#Xu2022Contrastive]_
+===================  ===================  ==================  =====  ===========  ========================================
 
 **(ii) Utility functions** :
 
@@ -185,6 +190,7 @@ Type                 Name                    Function
 Metric               eval_precision_at_k     Calculating Precision@k             `eval_precision_at_k <https://docs.pygod.org/en/latest/pygod.utils.html#pygod.utils.metric.eval_precision_at_k>`_
 Metric               eval_recall_at_k        Calculating Recall@k                `eval_recall_at_k <https://docs.pygod.org/en/latest/pygod.utils.html#pygod.utils.metric.eval_recall_at_k>`_
 Metric               eval_roc_auc            Calculating ROC-AUC Score           `eval_roc_auc <https://docs.pygod.org/en/latest/pygod.utils.html#pygod.utils.metric.eval_roc_auc>`_
+Metric               eval_average_precision  Calculating average precision       `eval_average_precision <https://docs.pygod.org/en/latest/pygod.utils.html#pygod.utils.metric.eval_average_precision>`_
 Data                 gen_structure_outliers  Generating structural outliers      `gen_structure_outliers <https://docs.pygod.org/en/latest/pygod.utils.html#pygod.utils.outlier_generator.gen_structure_outliers>`_
 Data                 gen_attribute_outliers  Generating attribute outliers       `gen_attribute_outliers <https://docs.pygod.org/en/latest/pygod.utils.html#pygod.utils.outlier_generator.gen_attribute_outliers>`_
 ===================  ======================  ==================================  ======================================================================================================================================
@@ -255,4 +261,10 @@ Reference
 
 .. [#Wang2021One] Wang, X., Jin, B., Du, Y., Cui, P., Tan, Y. and Yang, Y., 2021. One-class graph neural networks for anomaly detection in attributed networks. Neural computing and applications.
 
+.. [#Liu2021Anomaly] Liu, Y., Li, Z., Pan, S., Gong, C., Zhou, C. and Karypis, G., 2021. Anomaly detection on attributed networks via contrastive self-supervised learning. IEEE transactions on neural networks and learning systems (TNNLS).
+
+.. [#Jin2021ANEMONE] Jin, M., Liu, Y., Zheng, Y., Chi, L., Li, Y. and Pan, S., 2021. ANEMONE: Graph Anomaly Detection with Multi-Scale Contrastive Learning. In Proceedings of the 30th ACM International Conference on Information & Knowledge Management (CIKM).
+
 .. [#Yuan2021Higher] Yuan, X., Zhou, N., Yu, S., Huang, H., Chen, Z. and Xia, F., 2021, December. Higher-order Structure Based Anomaly Detection on Attributed Networks. In 2021 IEEE International Conference on Big Data (Big Data).
+
+.. [#Xu2022Contrastive] Xu, Z., Huang, X., Zhao, Y., Dong, Y., and Li, J., 2022. Contrastive Attributed Network Anomaly Detection with Data Augmentation. In Proceedings of the 26th Pacific-Asia Conference on Knowledge Discovery and Data Mining (PAKDD).
diff --git a/docs/api_cc.rst b/docs/api_cc.rst
index 362b499..4ac7fbe 100644
--- a/docs/api_cc.rst
+++ b/docs/api_cc.rst
@@ -5,11 +5,6 @@ The following APIs are applicable for all detector models for easy use.
 
 * :func:`pygod.models.base.BaseDetector.fit`: Fit detector. y is ignored in unsupervised methods.
 * :func:`pygod.models.base.BaseDetector.decision_function`: Predict raw anomaly scores of PyG Graph G using the fitted detector
-* :func:`pygod.models.base.BaseDetector.predict`: Predict if a particular sample is an outlier or not using the fitted detector.
-* :func:`pygod.models.base.BaseDetector.predict_proba`: Predict the probability of a sample being outlier using the fitted detector.
-* :func:`pygod.models.base.BaseDetector.predict_confidence`: Predict the model's sample-wise confidence (available in predict and predict_proba).
-* :func:`pygod.models.base.BaseDetector.process_graph` (you do not need to call this explicitly): Process the raw PyG data object into a tuple of sub data objects needed for the underlying model.
-
 
 Key Attributes of a fitted model:
 
@@ -17,10 +12,18 @@ Key Attributes of a fitted model:
   Outliers tend to have higher scores.
 * :attr:`pygod.models.base.BaseDetector.labels_`: The binary labels of the training data. 0 stands for inliers and 1 for outliers/anomalies.
 
+For the inductive setting:
+
+* :func:`pygod.models.base.BaseDetector.predict`: Predict if a particular sample is an outlier or not using the fitted detector.
+* :func:`pygod.models.base.BaseDetector.predict_proba`: Predict the probability of a sample being outlier using the fitted detector.
+* :func:`pygod.models.base.BaseDetector.predict_confidence`: Predict the model's sample-wise confidence (available in predict and predict_proba).
+
 
 **Input of PyGOD**: Please pass in a `PyTorch Geometric (PyG) <https://www.pyg.org/>`_ data object.
 See `PyG data processing examples <https://pytorch-geometric.readthedocs.io/en/latest/notes/introduction.html#data-handling-of-graphs>`_.
 
+* :func:`pygod.models.base.BaseDetector.process_graph` (you do not need to call this explicitly): Process the raw PyG data object into a tuple of sub data objects needed for the underlying model.
+
 
 See base class definition below:
 
diff --git a/docs/examples/intro.py b/docs/examples/intro.py
index 890f138..8921d1f 100644
--- a/docs/examples/intro.py
+++ b/docs/examples/intro.py
@@ -1,10 +1,10 @@
 """
-DOMINANT Example
+Model Example
 ================
 
 In this introductory tutorial, you will learn the basic workflow of
-PyGOD. This tutorial assumes that you have basic familiarity with
-PyTorch and PyTorch Geometric (PyG).
+PyGOD with a model example of DOMINANT. This tutorial assumes that
+you have basic familiarity with PyTorch and PyTorch Geometric (PyG).
 
 (Time estimate: 5 minutes)
 """
@@ -63,8 +63,9 @@
 #######################################################################
 # Inference
 # ---------
-# Then, your model is ready to use. We provide several inference methods.
-#######################################################################
+# Then, your model is ready to use.
+# We provide several inference methods.
+#
 # To predict the labels only:
 
 
diff --git a/docs/index.rst b/docs/index.rst
index bd66e96..a400be9 100644
--- a/docs/index.rst
+++ b/docs/index.rst
@@ -47,18 +47,17 @@ suspicious activities in social networks :cite:`dou2020enhancing`  and security
 
 PyGOD includes more than **10** latest graph-based detection algorithms,
 such as Dominant (SDM'19) and GUIDE (BigData'21).
-For consistently and accessibility, PyGOD is developed on top of `PyTorch Geometric (PyG) <https://www.pyg.org/>`_
+For consistency and accessibility, PyGOD is developed on top of `PyTorch Geometric (PyG) <https://www.pyg.org/>`_
 and `PyTorch <https://pytorch.org/>`_, and follows the API design of `PyOD <https://github.com/yzhao062/pyod>`_.
 See examples below for detecting anomalies with PyGOD in 5 lines!
 
-**PyGOD** is under actively developed and will be updated frequently!
-Please **star**, **watch**, and **fork**.
 
 **PyGOD is featured for**:
 
 * **Unified APIs, detailed documentation, and interactive examples** across various graph-based algorithms.
 * **Comprehensive coverage** of more than 10 latest graph outlier detectors.
-* **Full support of detections at multiple levels**, such as node-, edge-, and graph-level tasks (WIP).
+* **Full support of detections at multiple levels**, such as node-, edge- (WIP), and graph-level tasks (WIP).
+* **Scalable design for processing large graphs** via mini-batch and sampling.
 * **Streamline data processing with PyG**--fully compatible with PyG data objects.
 
 **Outlier Detection Using PyGOD with 5 Lines of Code**\ :
@@ -70,37 +69,35 @@ Please **star**, **watch**, and **fork**.
     # train a dominant detector
     from pygod.models import DOMINANT
 
-    model = DOMINANT()  # hyperparameters can be set here
+    model = DOMINANT(num_layers=4, epoch=20)  # hyperparameters can be set here
     model.fit(data)  # data is a Pytorch Geometric data object
 
     # get outlier scores on the input data
     outlier_scores = model.decision_scores # raw outlier scores on the input data
 
-    # predict on the new data
+    # predict on the new data in the inductive setting
     outlier_scores = model.decision_function(test_data) # raw outlier scores on the input data  # predict raw outlier scores on test
 
 
 
-**Citing PyGOD (to be announced soon)**\ :
+**Citing PyGOD**\ :
 
-`PyGOD paper <https://pygod.org>`_ will be available on arxiv soon.
+`PyGOD paper <https://arxiv.org/abs/2204.12095>`_ is available on arxiv :cite:`pygod2022`.
 If you use PyGOD in a scientific publication, we would appreciate
-citations to the following paper (to be announced)::
+citations to the following paper::
 
-    @article{tba,
-      author  = {tba},
-      title   = {PyGOD: A Comprehensive Python Library for Graph Outlier Detection},
-      journal = {tba},
+    @article{pygod2022,
+      author  = {Liu, Kay and Dou, Yingtong and Zhao, Yue and Ding, Xueying and Hu, Xiyang and Zhang, Ruitong and Ding, Kaize and Chen, Canyu and Peng, Hao and Shu, Kai and Chen, George H. and Jia, Zhihao and Yu, Philip S.},
+      title   = {PyGOD: A Python Library for Graph Outlier Detection},
+      journal = {arXiv preprint arXiv:2204.12095},
       year    = {2022},
     }
 
 or::
 
-    tba, 2022. PyGOD: A Comprehensive Python Library for Graph Outlier Detection. tba.
+    Liu, K., Dou, Y., Zhao, Y., Ding, X., Hu, X., Zhang, R., Ding, K., Chen, C., Peng, H., Shu, K., Chen, G.H., Jia, Z., and Yu, P.S. 2022. PyGOD: A Python Library for Graph Outlier Detection. arXiv preprint arXiv:2204.12095.
 
 
-----
-
 
 Implemented Algorithms
 ======================
@@ -109,20 +106,23 @@ PyGOD toolkit consists of two major functional groups:
 
 **(i) Node-level detection** :
 
-===================  ===================  ==================  ======================================================================================================  =====  ==============================================
-Type                 Backbone             Abbr                Algorithm                                                                                               Year   Class
-===================  ===================  ==================  ======================================================================================================  =====  ==============================================
-Unsupervised         NN                   MLPAE               Neural Networks and Deep Learning                                                                       2021   :class:`pygod.models.mlpae.MLPAE`
-Unsupervised         GNN                  GCNAE               Variational Graph Auto-Encoders                                                                         2021   :class:`pygod.models.gcnae.GCNAE`
-Unsupervised         MF                   ONE                 Outlier aware network embedding for attributed networks                                                 2019   :class:`pygod.models.one.ONE`
-Unsupervised         GNN                  DOMINANT            Deep anomaly detection on attributed networks                                                           2019   :class:`pygod.models.dominant.DOMINANT`
-Unsupervised         GNN                  DONE                Outlier Resistant Unsupervised Deep Architectures for Attributed Network Embedding                      2020   :class:`pygod.models.done.DONE`
-Unsupervised         GNN                  AdONE               Outlier Resistant Unsupervised Deep Architectures for Attributed Network Embedding                      2020   :class:`pygod.models.adone.AdONE`
-Unsupervised         GNN                  AnomalyDAE          AnomalyDAE: Dual autoencoder for anomaly detection on attributed networks                               2020   :class:`pygod.models.anomalydae.AnomalyDAE`
-Unsupervised         GAN                  GAAN                Generative Adversarial Attributed Network Anomaly Detection                                             2020   :class:`pygod.models.gaan.GAAN`
-Unsupervised         GNN                  OCGNN               One-Class Graph Neural Networks for Anomaly Detection in Attributed Networks                            2021   :class:`pygod.models.ocgnn.OCGNN`
-Unsupervised         GNN                  GUIDE               Higher-order Structure Based Anomaly Detection on Attributed Networks                                   2021   :class:`pygod.models.guide.GUIDE`
-===================  ===================  ==================  ======================================================================================================  =====  ==============================================
+===================  ===================  ==================  =====  ===========  ==============================================
+Type                 Backbone             Abbr                Year   Sampling     Class
+===================  ===================  ==================  =====  ===========  ==============================================
+Unsupervised         NN                   MLPAE               2014   Yes          :class:`pygod.models.MLPAE`
+Unsupervised         GNN                  GCNAE               2016   Yes          :class:`pygod.models.GCNAE`
+Unsupervised         MF                   ONE                 2019   No           :class:`pygod.models.ONE`
+Unsupervised         GNN                  DOMINANT            2019   Yes          :class:`pygod.models.DOMINANT`
+Unsupervised         GNN                  DONE                2020   Yes          :class:`pygod.models.DONE`
+Unsupervised         GNN                  AdONE               2020   Yes          :class:`pygod.models.AdONE`
+Unsupervised         GNN                  AnomalyDAE          2020   Yes          :class:`pygod.models.AnomalyDAE`
+Unsupervised         GAN                  GAAN                2020   Yes          :class:`pygod.models.GAAN`
+Unsupervised         GNN                  OCGNN               2021   Yes          :class:`pygod.models.OCGNN`
+Unsupervised/SSL     GNN                  CoLA (beta)         2021   In progress  :class:`pygod.models.CoLA`
+Unsupervised/SSL     GNN                  ANEMONE (beta)      2021   In progress  :class:`pygod.models.ANEMONE`
+Unsupervised         GNN                  GUIDE               2021   Yes          :class:`pygod.models.GUIDE`
+Unsupervised/SSL     GNN                  CONAD               2022   Yes          :class:`pygod.models.CONAD`
+===================  ===================  ==================  =====  ===========  ==============================================
 
 
 **(ii) Utility functions** :
@@ -133,6 +133,7 @@ Type                 Name                    Function
 Metric               eval_precision_at_k     Calculating Precision@k             `eval_precision_at_k <https://docs.pygod.org/en/latest/pygod.utils.html#pygod.utils.metric.eval_precision_at_k>`_
 Metric               eval_recall_at_k        Calculating Recall@k                `eval_recall_at_k <https://docs.pygod.org/en/latest/pygod.utils.html#pygod.utils.metric.eval_recall_at_k>`_
 Metric               eval_roc_auc            Calculating ROC-AUC Score           `eval_roc_auc <https://docs.pygod.org/en/latest/pygod.utils.html#pygod.utils.metric.eval_roc_auc>`_
+Metric               eval_average_precision  Calculating average precision       `eval_average_precision <https://docs.pygod.org/en/latest/pygod.utils.html#pygod.utils.metric.eval_average_precision>`_
 Data                 gen_structure_outliers  Generating structural outliers      `gen_structure_outliers <https://docs.pygod.org/en/latest/pygod.utils.html#pygod.utils.outlier_generator.gen_structure_outliers>`_
 Data                 gen_attribute_outliers  Generating attribute outliers       `gen_attribute_outliers <https://docs.pygod.org/en/latest/pygod.utils.html#pygod.utils.outlier_generator.gen_attribute_outliers>`_
 ===================  ======================  ==================================  ======================================================================================================================================
diff --git a/docs/install.rst b/docs/install.rst
index 9cb1599..60fe941 100644
--- a/docs/install.rst
+++ b/docs/install.rst
@@ -22,11 +22,9 @@ Alternatively, you could clone and run setup.py file:
 **Required Dependencies**\ :
 
 * Python 3.6 +
-* argparse>=1.4.0
 * numpy>=1.19.4
 * scikit-learn>=0.22.1
 * scipy>=1.5.2
-* pandas>=1.1.3
 * setuptools>=50.3.1.post20201107
 
 
diff --git a/docs/pygod.models.rst b/docs/pygod.models.rst
index fdd120d..230c748 100644
--- a/docs/pygod.models.rst
+++ b/docs/pygod.models.rst
@@ -1,113 +1,130 @@
 Models
 ==========
 
-pygod.models.adone module
-----------------------------
+AdONE
+-----
 
-.. automodule:: pygod.models.adone
-    :members: AdONE
-    :exclude-members: AdONE_Base, NeighDiff
+.. autoclass:: pygod.models.AdONE
+    :members:
+    :exclude-members: loss_func, process_graph
     :undoc-members:
     :show-inheritance:
     :inherited-members:
 
-pygod.models.anomalydae module
-------------------------------
+AnomalyDAE
+----------
 
-.. automodule:: pygod.models.anomalydae
-    :members: AnomalyDAE
-    :exclude-members: StructureAE, AttributeAE, AnomalyDAE_Base
+.. autoclass:: pygod.models.AnomalyDAE
+    :members:
+    :exclude-members: loss_func, process_graph
     :undoc-members:
     :show-inheritance:
     :inherited-members:
 
+CoLA
+----
 
+.. autoclass:: pygod.models.CoLA
+    :members:
+    :exclude-members: loss_func, process_graph
+    :undoc-members:
+    :show-inheritance:
+    :inherited-members:
 
-pygod.models.dominant module
-----------------------------
+CONAD
+-----
 
-.. automodule:: pygod.models.dominant
-    :members: DOMINANT
-    :exclude-members: DOMINANT_Base
+.. autoclass:: pygod.models.CONAD
+    :members:
+    :exclude-members: cal_outlierScore, calc_lossValues, process_graph, loss_func
     :undoc-members:
     :show-inheritance:
     :inherited-members:
 
-pygod.models.done module
-----------------------------
+DOMINANT
+--------
 
-.. automodule:: pygod.models.done
-    :members: DONE
-    :exclude-members: DONE_Base, NeighDiff
+.. autoclass:: pygod.models.DOMINANT
+    :members:
+    :exclude-members: loss_func, process_graph
     :undoc-members:
     :show-inheritance:
     :inherited-members:
 
+DONE
+----
 
-pygod.models.gaan module
-------------------------
-
-.. automodule:: pygod.models.gaan
-    :members: GAAN
-    :exclude-members: Generator, Discriminator, Encoder
+.. autoclass:: pygod.models.DONE
+    :members:
+    :exclude-members: loss_func, process_graph
     :undoc-members:
     :show-inheritance:
     :inherited-members:
 
-pygod.models.gcnae module
--------------------------
 
-.. automodule:: pygod.models.gcnae
-    :members: GCNAE
-    :exclude-members: Generator, Discriminator, Encoder
+GAAN
+----
+
+.. autoclass:: pygod.models.GAAN
+    :members:
+    :exclude-members: process_graph, score_function, train_model, loss_function
     :undoc-members:
     :show-inheritance:
     :inherited-members:
 
+GCNAE
+-----
 
-pygod.models.guide module
--------------------------
-
-.. automodule:: pygod.models.guide
-    :members: GUIDE
-    :exclude-members: GUIDE_Base, GNA, GNAConv
+.. autoclass:: pygod.models.GCNAE
+    :members:
+    :exclude-members: loss_func, process_graph
     :undoc-members:
     :show-inheritance:
     :inherited-members:
 
-pygod.models.mlpae module
--------------------------
 
-.. automodule:: pygod.models.mlpae
-    :members: MLPAE
-    :exclude-members: GUIDE_Base, GNA, GNAConv
+GUIDE
+-----
+
+.. autoclass:: pygod.models.GUIDE
+    :members:
+    :exclude-members: loss_func, process_graph
     :undoc-members:
     :show-inheritance:
     :inherited-members:
 
-pygod.models.ocgnn module
--------------------------
+MLPAE
+-----
 
-.. automodule:: pygod.models.ocgnn
-    :members: OCGNN
-    :exclude-members: GCN_base
+.. autoclass:: pygod.models.MLPAE
+    :members:
+    :exclude-members: loss_func, process_graph
     :undoc-members:
     :show-inheritance:
     :inherited-members:
 
-pygod.models.one module
------------------------
+OCGNN
+-----
 
-.. automodule:: pygod.models.one
-    :members: ONE
-    :exclude-members:
+.. autoclass:: pygod.models.OCGNN
+    :members:
+    :exclude-members: anomaly_scores, get_radius, process_graph, init_center, loss_function
     :undoc-members:
     :show-inheritance:
     :inherited-members:
 
+ONE
+---
 
+.. autoclass:: pygod.models.ONE
+    :members:
+    :exclude-members: cal_outlierScore, calc_lossValues, process_graph, loss_func
+    :undoc-members:
+    :show-inheritance:
+    :inherited-members:
 
-.. rubric:: References
+Reference
+---------
 
 .. bibliography::
    :cited:
diff --git a/docs/pygod.utils.rst b/docs/pygod.utils.rst
index 0381e2b..41af3b9 100644
--- a/docs/pygod.utils.rst
+++ b/docs/pygod.utils.rst
@@ -2,32 +2,30 @@ Utility Functions
 ==================
 
 
-pygod.utils.early\_stopping module
-----------------------------------
+Early Stopping
+--------------
 
-.. automodule:: pygod.utils.early_stopping
+.. autoclass:: pygod.utils.EarlyStopping
     :members:
     :undoc-members:
-    :show-inheritance:
-    :inherited-members:
 
 
-pygod.utils.metric module
--------------------------
+Metric
+------
 
 .. automodule:: pygod.utils.metric
     :members:
     :undoc-members:
-    :show-inheritance:
-    :inherited-members:
 
 
-pygod.utils.outlier\_generator module
--------------------------------------
+Outlier Generator
+-----------------
 
-.. automodule:: pygod.utils.outlier_generator
+.. autoclass:: pygod.utils.gen_attribute_outliers
     :members:
     :undoc-members:
-    :show-inheritance:
     :inherited-members:
 
+.. autoclass:: pygod.utils.gen_structure_outliers
+    :members:
+    :undoc-members:
diff --git a/docs/requirements.txt b/docs/requirements.txt
index e97ed02..719043b 100644
--- a/docs/requirements.txt
+++ b/docs/requirements.txt
@@ -1,12 +1,11 @@
-argparse>=1.4.0
 furo
 networkx>=2.6.3
 numpy>=1.19.4
-pandas>=1.1.3
 https://download.pytorch.org/whl/cpu/torch-1.10.0%2Bcpu-cp38-cp38-linux_x86_64.whl
 https://data.pyg.org/whl/torch-1.10.0%2Bcpu/torch_scatter-2.0.9-cp38-cp38-linux_x86_64.whl
 https://data.pyg.org/whl/torch-1.10.0%2Bcpu/torch_sparse-0.6.12-cp38-cp38-linux_x86_64.whl
-torch_geometric>=2.0.3
+https://data.pyg.org/whl/torch-1.10.0%2Bcpu/torch_cluster-1.6.0-cp38-cp38-linux_x86_64.whl
+torch_geometric>=2.0.4
 scipy>=1.5.2
 setuptools>=50.3.1.post20201107
 scikit-learn>=0.22.1
diff --git a/docs/zreferences.bib b/docs/zreferences.bib
index 39281b6..024815a 100644
--- a/docs/zreferences.bib
+++ b/docs/zreferences.bib
@@ -111,4 +111,26 @@ @inproceedings{cai2021structural
   booktitle={Proceedings of the 30th ACM International Conference on Information \& Knowledge Management},
   pages={3747--3756},
   year={2021}
-}
\ No newline at end of file
+}
+
+@inproceedings{xu2022contrastive,
+  title = {Contrastive Attributed Network Anomaly Detection with Data Augmentation},
+  author = {Xu, Zhiming and Huang, Xiao and Zhao, Yue and Dong, Yushun and Li, Jundong},
+  booktitle={Pacific-Asian Conference on Knowledge Discovery and Data Mining (PAKDD)},
+  year = {2022}
+}
+
+@article{pygod2022,
+  author  = {Liu, Kay and Dou, Yingtong and Zhao, Yue and Ding, Xueying and Hu, Xiyang and Zhang, Ruitong and Ding, Kaize and Chen, Canyu and Peng, Hao and Shu, Kai and Chen, George H. and Jia, Zhihao and Yu, Philip S.},
+  title   = {PyGOD: A Python Library for Graph Outlier Detection},
+  journal = {arXiv preprint arXiv:2204.12095},
+  year    = {2022},
+}
+
+@inproceedings{jin2021anemone,
+  title={ANEMONE: Graph Anomaly Detection with Multi-Scale Contrastive Learning},
+  author={Jin, Ming and Liu, Yixin and Zheng, Yu and Chi, Lianhua and Li, Yuan-Fang and Pan, Shirui},
+  booktitle={Proceedings of the 30th ACM International Conference on Information \& Knowledge Management},
+  pages={3122--3126},
+  year={2021}
+}
diff --git a/pygod/__init__.py b/pygod/__init__.py
index af38f37..07ac336 100644
--- a/pygod/__init__.py
+++ b/pygod/__init__.py
@@ -1,2 +1,3 @@
 from . import models
 from . import utils
+from .version import __version__
diff --git a/pygod/models/__init__.py b/pygod/models/__init__.py
index 31e0bd6..482ddb9 100644
--- a/pygod/models/__init__.py
+++ b/pygod/models/__init__.py
@@ -9,7 +9,9 @@
 from .mlpae import MLPAE
 from .ocgnn import OCGNN
 from .one import ONE
-
+from .conad import CONAD
+from .cola import CoLA
+from .anemone import ANEMONE
 
 
 
diff --git a/pygod/models/adone.py b/pygod/models/adone.py
index f3308bf..8d8aa24 100644
--- a/pygod/models/adone.py
+++ b/pygod/models/adone.py
@@ -4,21 +4,24 @@
 # License: BSD 2 clause
 
 import torch
+import numpy as np
 import torch.nn as nn
 import torch.nn.functional as F
-from torch_geometric.utils import to_dense_adj
 from torch_geometric.nn import MessagePassing
+from torch_geometric.utils import to_dense_adj
+from torch_geometric.loader import NeighborLoader
 from sklearn.utils.validation import check_is_fitted
 
 from . import BaseDetector
 from .basic_nn import MLP
+from ..utils.utility import validate_device
 from ..utils.metric import eval_roc_auc
 
 
 class AdONE(BaseDetector):
     """
     AdONE (Adversarial Outlier Aware Attributed Network Embedding)
-    AdONE is consist of an attribute autoencoder and a structure
+    AdONE consist of an attribute autoencoder and a structure
     autoencoder. It estimates five loss to optimize the model,
     including an attribute proximity loss, an attribute homophily loss,
     a structure proximity loss, a structure homophily loss, and an
@@ -68,6 +71,11 @@ class AdONE(BaseDetector):
         Maximum number of training epoch. Default: ``5``.
     gpu : int
         GPU Index, -1 for using CPU. Default: ``0``.
+    batch_size : int, optional
+        Minibatch size, 0 for full batch training. Default: ``0``.
+    num_neigh : int, optional
+        Number of neighbors in sampling, -1 for all neighbors.
+        Default: ``-1``.
     verbose : bool
         Verbosity mode. Turn on to print out log information.
         Default: ``False``.
@@ -95,6 +103,8 @@ def __init__(self,
                  lr=5e-3,
                  epoch=5,
                  gpu=0,
+                 batch_size=0,
+                 num_neigh=-1,
                  verbose=False):
         super(AdONE, self).__init__(contamination=contamination)
 
@@ -113,10 +123,9 @@ def __init__(self,
         # training param
         self.lr = lr
         self.epoch = epoch
-        if gpu >= 0 and torch.cuda.is_available():
-            self.device = 'cuda:{}'.format(gpu)
-        else:
-            self.device = 'cpu'
+        self.device = validate_device(gpu)
+        self.batch_size = batch_size
+        self.num_neigh = num_neigh
 
         # other param
         self.verbose = verbose
@@ -124,29 +133,32 @@ def __init__(self,
 
     def fit(self, G, y_true=None):
         """
-        Description
-        -----------
         Fit detector with input data.
 
         Parameters
         ----------
-        G : PyTorch Geometric Data instance (torch_geometric.data.Data)
+        G : torch_geometric.data.Data
             The input data.
-        y_true : numpy.array, optional (default=None)
-            The optional outlier ground truth labels used to monitor the
-            training progress. They are not used to optimize the
-            unsupervised model.
+        y_true : numpy.ndarray, optional
+            The optional outlier ground truth labels used to monitor
+            the training progress. They are not used to optimize the
+            unsupervised model. Default: ``None``.
 
         Returns
         -------
         self : object
             Fitted estimator.
         """
-
-        x, s, edge_index = self.process_graph(G)
-
-        self.model = AdONE_Base(x_dim=x.shape[1],
-                                s_dim=s.shape[1],
+        G.node_idx = torch.arange(G.x.shape[0])
+        G.s = to_dense_adj(G.edge_index)[0]
+        if self.batch_size == 0:
+            self.batch_size = G.x.shape[0]
+        loader = NeighborLoader(G,
+                                [self.num_neigh],
+                                batch_size=self.batch_size)
+
+        self.model = AdONE_Base(x_dim=G.x.shape[1],
+                                s_dim=G.s.shape[1],
                                 hid_dim=self.hid_dim,
                                 num_layers=self.num_layers,
                                 dropout=self.dropout,
@@ -156,34 +168,48 @@ def fit(self, G, y_true=None):
                                      lr=self.lr,
                                      weight_decay=self.weight_decay)
 
-        score = None
+        self.model.train()
+        decision_scores = np.zeros(G.x.shape[0])
         for epoch in range(self.epoch):
-            self.model.train()
-            x_, s_, h_a, h_s, dna, dns, dis_a, dis_s \
-                = self.model(x, s, edge_index)
-            score, loss = self.loss_func(x, x_, s, s_, h_a, h_s,
-                                         dna, dns, dis_a, dis_s)
-
-            optimizer.zero_grad()
-            loss.backward()
-            optimizer.step()
+            epoch_loss = 0
+            for sampled_data in loader:
+                batch_size = sampled_data.batch_size
+                node_idx = sampled_data.node_idx
+                x, s, edge_index = self.process_graph(G)
+                x_, s_, h_a, h_s, dna, dns, dis_a, dis_s \
+                    = self.model(x, s, edge_index)
+                score, loss = self.loss_func(x[:batch_size],
+                                             x_[:batch_size],
+                                             s[:batch_size],
+                                             s_[:batch_size],
+                                             h_a[:batch_size],
+                                             h_s[:batch_size],
+                                             dna[:batch_size],
+                                             dns[:batch_size],
+                                             dis_a[:batch_size],
+                                             dis_s[:batch_size])
+                epoch_loss += loss.item() * batch_size
+                decision_scores[node_idx[:batch_size]] = score.detach() \
+                                                              .cpu().numpy()
+
+                optimizer.zero_grad()
+                loss.backward()
+                optimizer.step()
 
             if self.verbose:
                 print("Epoch {:04d}: Loss {:.4f}"
-                      .format(epoch, loss.item()), end='')
+                      .format(epoch, epoch_loss / G.x.shape[0]), end='')
                 if y_true is not None:
-                    auc = eval_roc_auc(y_true, score.detach().cpu().numpy())
+                    auc = eval_roc_auc(y_true, decision_scores)
                     print(" | AUC {:.4f}".format(auc), end='')
                 print()
 
-        self.decision_scores_ = score.detach().cpu().numpy()
+        self.decision_scores_ = decision_scores
         self._process_decision_scores()
         return self
 
     def decision_function(self, G):
         """
-        Description
-        -----------
         Predict raw anomaly score using the fitted detector. Outliers
         are assigned with larger anomaly scores.
 
@@ -198,23 +224,41 @@ def decision_function(self, G):
             The anomaly score of shape :math:`N`.
         """
         check_is_fitted(self, ['model'])
+        G.node_idx = torch.arange(G.x.shape[0])
+        G.s = to_dense_adj(G.edge_index)[0]
+        if self.batch_size == 0:
+            self.batch_size = G.x.shape[0]
+        loader = NeighborLoader(G,
+                                [self.num_neigh],
+                                batch_size=self.batch_size)
 
-        # get needed data object from the input data
-        x, s, edge_index = self.process_graph(G)
-
-        # enable the evaluation mode
         self.model.eval()
+        outlier_scores = np.zeros(G.x.shape[0])
+        for sampled_data in loader:
+            batch_size = sampled_data.batch_size
+            node_idx = sampled_data.node_idx
 
-        # construct the vector for holding the reconstruction error
-        x_, s_, h_a, h_s, dna, dns, dis_a, dis_s = self.model(x, s, edge_index)
-        outlier_scores, _ = self.loss_func(x, x_, s, s_, h_a, h_s,
-                                           dna, dns, dis_a, dis_s)
-        return outlier_scores.detach().cpu().numpy()
+            x, s, edge_index = self.process_graph(G)
+
+            x_, s_, h_a, h_s, dna, dns, dis_a, dis_s \
+                = self.model(x, s, edge_index)
+            score, _ = self.loss_func(x[:batch_size],
+                                      x_[:batch_size],
+                                      s[:batch_size],
+                                      s_[:batch_size],
+                                      h_a[:batch_size],
+                                      h_s[:batch_size],
+                                      dna[:batch_size],
+                                      dns[:batch_size],
+                                      dis_a[:batch_size],
+                                      dis_s[:batch_size])
+
+            outlier_scores[node_idx[:batch_size]] = score.detach() \
+                                                         .cpu().numpy()
+        return outlier_scores
 
     def process_graph(self, G):
         """
-        Description
-        -----------
         Process the raw PyG data object into a tuple of sub data
         objects needed for the model.
 
diff --git a/pygod/models/anemone.py b/pygod/models/anemone.py
new file mode 100644
index 0000000..ab47128
--- /dev/null
+++ b/pygod/models/anemone.py
@@ -0,0 +1,717 @@
+# -*- coding: utf-8 -*-
+"""ANEMONE: Graph Anomaly Detection 
+with Multi-Scale Contrastive Learning (ANEMONE)"""
+# Author: Canyu Chen <cchen151@hawk.iit.edu>
+# License: BSD 2 clause
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+import scipy.sparse as sp
+import random
+import os
+
+from torch_geometric.utils import to_dense_adj
+from torch_cluster import random_walk
+
+from . import BaseDetector
+from ..utils.metric import eval_roc_auc
+
+os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"
+
+
+def sparse_to_tuple(sparse_mx, insert_batch=False):
+    """Convert sparse matrix to tuple representation."""
+    """Set insert_batch=True if you want to insert a batch dimension."""
+
+    def to_tuple(mx):
+        if not sp.isspmatrix_coo(mx):
+            mx = mx.tocoo()
+        if insert_batch:
+            coords = np.vstack((np.zeros(mx.row.shape[0]), mx.row, mx.col)).T
+            values = mx.data
+            shape = (1,) + mx.shape
+        else:
+            coords = np.vstack((mx.row, mx.col)).transpose()
+            values = mx.data
+            shape = mx.shape
+        return coords, values, shape
+
+    if isinstance(sparse_mx, list):
+        for i in range(len(sparse_mx)):
+            sparse_mx[i] = to_tuple(sparse_mx[i])
+    else:
+        sparse_mx = to_tuple(sparse_mx)
+
+    return sparse_mx
+
+
+def preprocess_features(features):
+    """Row-normalize feature matrix and convert to tuple representation"""
+    rowsum = np.array(features.sum(1))
+    r_inv = np.power(rowsum, -1).flatten()
+    r_inv[np.isinf(r_inv)] = 0.
+    r_mat_inv = sp.diags(r_inv)
+    features = r_mat_inv.dot(features)
+    return features.todense(), sparse_to_tuple(features)
+
+
+def normalize_adj(adj):
+    """Symmetrically normalize adjacency matrix."""
+    adj = sp.coo_matrix(adj)
+    rowsum = np.array(adj.sum(1))
+    d_inv_sqrt = np.power(rowsum, -0.5).flatten()
+    d_inv_sqrt[np.isinf(d_inv_sqrt)] = 0.
+    d_mat_inv_sqrt = sp.diags(d_inv_sqrt)
+    return adj.dot(d_mat_inv_sqrt).transpose().dot(d_mat_inv_sqrt).tocoo()
+
+
+def dense_to_one_hot(labels_dense, num_classes):
+    """Convert class labels from scalars to one-hot vectors."""
+    num_labels = labels_dense.shape[0]
+    index_offset = np.arange(num_labels) * num_classes
+    labels_one_hot = np.zeros((num_labels, num_classes))
+    labels_one_hot.flat[index_offset + labels_dense.ravel()] = 1
+    return labels_one_hot
+
+
+def generate_rw_subgraph(pyg_graph, nb_nodes, subgraph_size):
+    """Generate subgraph with random walk algorithm."""
+    row, col = pyg_graph.edge_index
+    all_idx = torch.tensor(list(range(nb_nodes)))
+    traces = random_walk(row, col, all_idx, walk_length=3)
+    subv = traces.tolist()
+    return subv
+
+
+class ANEMONE(BaseDetector):
+    """
+    ANEMONE (ANEMONE: Graph Anomaly Detection 
+    with Multi-Scale Contrastive Learning)
+    ANEMONE is a multi-scale contrastive self-supervised 
+    learning-based method for graph anomaly detection. (beta)
+
+    Parameters
+    ----------
+    lr : float, optional
+        Learning rate. 
+    verbose : bool
+        Verbosity mode. Turn on to print out log information.
+        Default: ``False``.
+    epoch : int, optional
+        Maximum number of training epoch. 
+    embedding_dim : 
+    negsamp_ratio : 
+    readout : 
+    dataset : 
+    weight_decay : 
+    batch_size : 
+    subgraph_size : 
+    alpha :
+    negsamp_ratio_patch :
+    negsamp_ratio_context :
+    auc_test_rounds : 
+    contamination : float, optional
+        Valid in (0., 0.5). The proportion of outliers in the data set.
+        Used when fitting to define the threshold on the decision
+        function. Default: ``0.1``.
+    gpu : int
+        GPU Index, -1 for using CPU. Default: ``0``.
+
+    Examples
+    --------
+    >>> from pygod.models import ANEMONE
+    >>> model = ANEMONE()
+    >>> model.fit(data)
+    >>> prediction = model.predict(data)
+    """
+
+    def __init__(self,
+                 lr=None,
+                 verbose=False,
+                 epoch=None,
+                 embedding_dim=64,
+                 negsamp_ratio=1,
+                 readout='avg',
+                 dataset='cora',
+                 weight_decay=0.0,
+                 batch_size=300,
+                 subgraph_size=4,
+                 alpha=1.0,
+                 negsamp_ratio_patch=1,
+                 negsamp_ratio_context=1,
+                 auc_test_rounds=256,
+                 contamination=0.1,
+                 gpu=0):
+        super(ANEMONE, self).__init__(contamination=contamination)
+
+        self.dataset = dataset
+        self.embedding_dim = embedding_dim
+        self.negsamp_ratio = negsamp_ratio
+        self.readout = readout
+        self.weight_decay = weight_decay
+        self.batch_size = batch_size
+        self.subgraph_size = subgraph_size
+        self.auc_test_rounds = auc_test_rounds
+        self.alpha = alpha
+        self.negsamp_ratio_patch = negsamp_ratio_patch
+        self.negsamp_ratio_context = negsamp_ratio_context
+        if gpu >= 0 and torch.cuda.is_available():
+            self.device = 'cuda:{}'.format(gpu)
+        else:
+            self.device = 'cpu'
+
+        if lr is None:
+            if self.dataset in ['cora', 'citeseer', 'pubmed', 'Flickr']:
+                self.lr = 1e-3
+            elif self.dataset == 'ACM':
+                self.lr = 5e-4
+            elif self.dataset == 'BlogCatalog':
+                self.lr = 3e-3
+            else:
+                self.lr = 1e-3
+
+        if epoch is None:
+            if self.dataset in ['cora', 'citeseer', 'pubmed']:
+                self.num_epoch = 100
+            elif self.dataset in ['BlogCatalog', 'Flickr', 'ACM']:
+                self.num_epoch = 400
+            else:
+                self.num_epoch = 100
+
+        self.verbose = verbose
+        self.model = None
+
+    def fit(self, G):
+        """
+        Fit detector with input data.
+
+        Parameters
+        ----------
+        G : PyTorch Geometric Data instance (torch_geometric.data.Data)
+            The input data.
+
+        Returns
+        -------
+        self : object
+            Fitted estimator.
+        """
+
+        x, adj, edge_index, labels = self.process_graph(G)
+
+        adj = adj.cpu().numpy()
+        x = x.cpu().numpy()
+
+        nb_nodes = x.shape[0]
+        ft_size = x.shape[1]
+
+        adj = normalize_adj(adj)
+        adj = (adj + sp.eye(adj.shape[0])).todense()
+
+        x = torch.FloatTensor(x[np.newaxis])
+        adj = torch.FloatTensor(adj[np.newaxis])
+
+        self.model = ANEMONE_Base(ft_size,
+                                  self.embedding_dim,
+                                  'prelu',
+                                  self.negsamp_ratio_patch,
+                                  self.negsamp_ratio_context,
+                                  self.readout)
+
+        optimiser = torch.optim.Adam(self.model.parameters(),
+                                     lr=self.lr,
+                                     weight_decay=self.weight_decay)
+
+        b_xent_patch = nn.BCEWithLogitsLoss(reduction='none',
+                                            pos_weight=torch.tensor(
+                                                [self.negsamp_ratio_patch]))
+        b_xent_context = nn.BCEWithLogitsLoss(reduction='none',
+                                              pos_weight=torch.tensor([
+                                                  self.negsamp_ratio_context]))
+
+        batch_num = nb_nodes // self.batch_size + 1
+
+        multi_epoch_ano_score = np.zeros((self.num_epoch, nb_nodes))
+
+        for epoch in range(self.num_epoch):
+
+            self.model.train()
+
+            all_idx = list(range(nb_nodes))
+            random.shuffle(all_idx)
+
+            subgraphs = generate_rw_subgraph(G, nb_nodes, self.subgraph_size)
+
+            for batch_idx in range(batch_num):
+
+                optimiser.zero_grad()
+
+                is_final_batch = (batch_idx == (batch_num - 1))
+                if not is_final_batch:
+                    idx = all_idx[batch_idx * self.batch_size:
+                                  (batch_idx + 1) * self.batch_size]
+                else:
+                    idx = all_idx[batch_idx * self.batch_size:]
+
+                cur_batch_size = len(idx)
+
+                lbl_patch = torch.unsqueeze(torch.cat(
+                    (torch.ones(cur_batch_size),
+                     torch.zeros(cur_batch_size * self.negsamp_ratio_patch))),
+                    1)
+
+                lbl_context = torch.unsqueeze(torch.cat(
+                    (torch.ones(cur_batch_size), torch.zeros(
+                        cur_batch_size * self.negsamp_ratio_context))), 1)
+
+                ba = []
+                bf = []
+                added_adj_zero_row = torch.zeros(
+                    (cur_batch_size, 1, self.subgraph_size))
+                added_adj_zero_col = torch.zeros(
+                    (cur_batch_size, self.subgraph_size + 1, 1))
+                added_adj_zero_col[:, -1, :] = 1.
+                added_feat_zero_row = torch.zeros((cur_batch_size, 1, ft_size))
+
+                for i in idx:
+                    cur_adj = adj[:, subgraphs[i], :][:, :, subgraphs[i]]
+                    cur_feat = x[:, subgraphs[i], :]
+                    ba.append(cur_adj)
+                    bf.append(cur_feat)
+
+                ba = torch.cat(ba)
+                ba = torch.cat((ba, added_adj_zero_row), dim=1)
+                ba = torch.cat((ba, added_adj_zero_col), dim=2)
+                bf = torch.cat(bf)
+                bf = torch.cat(
+                    (bf[:, :-1, :], added_feat_zero_row, bf[:, -1:, :]), dim=1)
+
+                logits_1, logits_2 = self.model(bf, ba)
+
+                # Context-level
+                loss_all_1 = b_xent_context(logits_1, lbl_context)
+                loss_1 = torch.mean(loss_all_1)
+
+                # Patch-level
+                loss_all_2 = b_xent_patch(logits_2, lbl_patch)
+                loss_2 = torch.mean(loss_all_2)
+
+                loss = self.alpha * loss_1 + (1 - self.alpha) * loss_2
+
+                loss.backward()
+                optimiser.step()
+
+                logits_1 = torch.sigmoid(torch.squeeze(logits_1))
+                logits_2 = torch.sigmoid(torch.squeeze(logits_2))
+
+                if self.alpha != 1.0 and self.alpha != 0.0:
+                    if self.negsamp_ratio_context == 1 and \
+                            self.negsamp_ratio_patch == 1:
+                        ano_score_1 = - (logits_1[:cur_batch_size] -
+                            logits_1[cur_batch_size:]).detach().cpu().numpy()
+                        ano_score_2 = - (logits_2[:cur_batch_size] -
+                            logits_2[cur_batch_size:]).detach().cpu().numpy()
+                    else:
+                        ano_score_1 = - (logits_1[:cur_batch_size] -
+                            torch.mean(logits_1[cur_batch_size:].view(
+                                cur_batch_size, self.negsamp_ratio_context),
+                                dim=1)).detach().cpu().numpy()  # context
+                        ano_score_2 = - (logits_2[:cur_batch_size] -
+                            torch.mean(logits_2[cur_batch_size:].view(
+                                    cur_batch_size, self.negsamp_ratio_patch),
+                                    dim=1)).detach().cpu().numpy()  # patch
+                    ano_score = self.alpha * ano_score_1 + (
+                                1 - self.alpha) * ano_score_2
+                elif self.alpha == 1.0:
+                    if self.negsamp_ratio_context == 1:
+                        ano_score = - (logits_1[:cur_batch_size] -
+                            logits_1[cur_batch_size:]).detach().cpu().numpy()
+                    else:
+                        ano_score = - (logits_1[:cur_batch_size] -
+                            torch.mean(logits_1[cur_batch_size:].view(
+                                cur_batch_size, self.negsamp_ratio_context),
+                                dim=1)).detach().cpu().numpy()  # context
+                elif self.alpha == 0.0:
+                    if self.negsamp_ratio_patch == 1:
+                        ano_score = - (logits_2[:cur_batch_size] -
+                            logits_2[cur_batch_size:]).detach().cpu().numpy()
+                    else:
+                        ano_score = - (logits_2[:cur_batch_size] -
+                            torch.mean(logits_2[cur_batch_size:].view(
+                                cur_batch_size, self.negsamp_ratio_patch),
+                                dim=1)).detach().cpu().numpy()  # patch
+
+                multi_epoch_ano_score[epoch, idx] = ano_score
+
+        ano_score_final = np.mean(multi_epoch_ano_score, axis=0)
+
+        self.decision_scores_ = ano_score_final
+        self._process_decision_scores()
+        return self
+
+    def decision_function(self, G):
+        """
+        Predict raw anomaly score using the fitted detector. Outliers
+        are assigned with larger anomaly scores.
+
+        Parameters
+        ----------
+        G : PyTorch Geometric Data instance (torch_geometric.data.Data)
+            The input data.
+
+        Returns
+        -------
+        outlier_scores : numpy.ndarray
+            The anomaly score of shape :math:`N`.
+        """
+
+        x, adj, edge_index, _ = self.process_graph(G)
+
+        adj = adj.cpu().numpy()
+        x = x.cpu().numpy()
+
+        nb_nodes = x.shape[0]
+        ft_size = x.shape[1]
+
+        adj = normalize_adj(adj)
+        adj = (adj + sp.eye(adj.shape[0])).todense()
+
+        x = torch.FloatTensor(x[np.newaxis])
+        adj = torch.FloatTensor(adj[np.newaxis])
+
+        batch_num = nb_nodes // self.batch_size + 1
+
+        multi_round_ano_score = np.zeros((self.auc_test_rounds, nb_nodes))
+
+        # enable the evaluation mode
+        self.model.eval()
+
+        for round in range(self.auc_test_rounds):
+
+            all_idx = list(range(nb_nodes))
+            random.shuffle(all_idx)
+
+            subgraphs = generate_rw_subgraph(G, nb_nodes, self.subgraph_size)
+
+            for batch_idx in range(batch_num):
+
+                is_final_batch = (batch_idx == (batch_num - 1))
+
+                if not is_final_batch:
+                    idx = all_idx[batch_idx * self.batch_size:
+                                  (batch_idx + 1) * self.batch_size]
+                else:
+                    idx = all_idx[batch_idx * self.batch_size:]
+
+                cur_batch_size = len(idx)
+
+                ba = []
+                bf = []
+                added_adj_zero_row = torch.zeros(
+                    (cur_batch_size, 1, self.subgraph_size))
+                added_adj_zero_col = torch.zeros(
+                    (cur_batch_size, self.subgraph_size + 1, 1))
+                added_adj_zero_col[:, -1, :] = 1.
+                added_feat_zero_row = torch.zeros((cur_batch_size, 1, ft_size))
+
+                for i in idx:
+                    cur_adj = adj[:, subgraphs[i], :][:, :, subgraphs[i]]
+                    cur_feat = x[:, subgraphs[i], :]
+                    ba.append(cur_adj)
+                    bf.append(cur_feat)
+
+                ba = torch.cat(ba)
+                ba = torch.cat((ba, added_adj_zero_row), dim=1)
+                ba = torch.cat((ba, added_adj_zero_col), dim=2)
+                bf = torch.cat(bf)
+                bf = torch.cat(
+                    (bf[:, :-1, :], added_feat_zero_row, bf[:, -1:, :]), dim=1)
+
+                with torch.no_grad():
+
+                    test_logits_1, test_logits_2 = self.model(bf, ba)
+                    test_logits_1 = torch.sigmoid(torch.squeeze(test_logits_1))
+                    test_logits_2 = torch.sigmoid(torch.squeeze(test_logits_2))
+
+                if self.alpha != 1.0 and self.alpha != 0.0:
+                    if self.negsamp_ratio_context == 1 and \
+                            self.negsamp_ratio_patch == 1:
+                        ano_score_1 = - (test_logits_1[:cur_batch_size] -
+                            test_logits_1[cur_batch_size:]).cpu().numpy()
+                        ano_score_2 = - (test_logits_2[:cur_batch_size] -
+                            test_logits_2[cur_batch_size:]).cpu().numpy()
+                    else:
+                        ano_score_1 = - (test_logits_1[:cur_batch_size] -
+                            torch.mean(test_logits_1[cur_batch_size:].view(
+                                cur_batch_size, self.negsamp_ratio_context),
+                                dim=1)).cpu().numpy()  # context
+                        ano_score_2 = - (test_logits_2[:cur_batch_size] -
+                            torch.mean(test_logits_2[cur_batch_size:].view(
+                                cur_batch_size, self.negsamp_ratio_patch),
+                                dim=1)).cpu().numpy()  # patch
+                    ano_score = self.alpha * ano_score_1 + \
+                                (1 - self.alpha) * ano_score_2
+                elif self.alpha == 1.0:
+                    if self.negsamp_ratio_context == 1:
+                        ano_score = - (test_logits_1[:cur_batch_size] -
+                            test_logits_1[cur_batch_size:]).cpu().numpy()
+                    else:
+                        ano_score = - (test_logits_1[:cur_batch_size] -
+                            torch.mean(test_logits_1[cur_batch_size:].view(
+                                cur_batch_size, self.negsamp_ratio_context),
+                                dim=1)).cpu().numpy()  # context
+                elif self.alpha == 0.0:
+                    if self.negsamp_ratio_patch == 1:
+                        ano_score = - (test_logits_2[:cur_batch_size] -
+                            test_logits_2[cur_batch_size:]).cpu().numpy()
+                    else:
+                        ano_score = - (test_logits_2[:cur_batch_size] -
+                            torch.mean(test_logits_2[cur_batch_size:].view(
+                                cur_batch_size, self.negsamp_ratio_patch),
+                                dim=1)).cpu().numpy()  # patch
+
+                multi_round_ano_score[round, idx] = ano_score
+
+        ano_score_final = np.mean(multi_round_ano_score, axis=0)
+
+        return ano_score_final
+
+    def process_graph(self, G):
+        """
+        Description
+        -----------
+        Process the raw PyG data object into a tuple of sub data
+        objects needed for the model.
+
+        Parameters
+        ----------
+        G : PyTorch Geometric Data instance (torch_geometric.data.Data)
+            The input data.
+
+        Returns
+        -------
+        x : torch.Tensor
+            Attribute (feature) of nodes.
+        adj : torch.Tensor
+            Adjacency matrix of the graph.
+        edge_index : torch.Tensor
+            Edge list of the graph.
+        y : torch.Tensor
+            Labels of nodes.
+        """
+        edge_index = G.edge_index
+
+        adj = to_dense_adj(edge_index)[0].to(self.device)
+
+        edge_index = edge_index.to(self.device)
+        adj = adj.to(self.device)
+        x = G.x.to(self.device)
+
+        if hasattr(G, 'y'):
+            y = G.y
+        else:
+            y = None
+
+        # return data objects needed for the network
+        return x, adj, edge_index, y
+
+
+class ANEMONE_Base(nn.Module):
+    def __init__(self,
+                 n_in,
+                 n_h,
+                 activation,
+                 negsamp_round_patch,
+                 negsamp_round_context,
+                 readout):
+        super(ANEMONE_Base, self).__init__()
+        self.read_mode = readout
+        self.gcn_context = GCN(n_in, n_h, activation)
+        self.gcn_patch = GCN(n_in, n_h, activation)
+        if readout == 'max':
+            self.read = MaxReadout()
+        elif readout == 'min':
+            self.read = MinReadout()
+        elif readout == 'avg':
+            self.read = AvgReadout()
+        elif readout == 'weighted_sum':
+            self.read = WSReadout()
+
+        self.c_disc = Contextual_Discriminator(n_h, negsamp_round_context)
+        self.p_disc = Patch_Discriminator(n_h, negsamp_round_patch)
+
+    def forward(self, seq1, adj, sparse=False, samp_bias1=None,
+                samp_bias2=None):
+        h_1 = self.gcn_context(seq1, adj, sparse)
+        h_2 = self.gcn_patch(seq1, adj, sparse)
+
+        if self.read_mode != 'weighted_sum':
+            c = self.read(h_1[:, :-1, :])
+            h_mv = h_1[:, -1, :]
+            h_unano = h_2[:, -1, :]
+            h_ano = h_2[:, -2, :]
+        else:
+            c = self.read(h_1[:, :-1, :], h_1[:, -2:-1, :])
+            h_mv = h_1[:, -1, :]
+            h_unano = h_2[:, -1, :]
+            h_ano = h_2[:, -2, :]
+
+        ret1 = self.c_disc(c, h_mv, samp_bias1, samp_bias2)
+        ret2 = self.p_disc(h_ano, h_unano, samp_bias1, samp_bias2)
+
+        return ret1, ret2
+
+
+class GCN(nn.Module):
+    def __init__(self, in_ft, out_ft, act, bias=True):
+        super(GCN, self).__init__()
+        self.fc = nn.Linear(in_ft, out_ft, bias=False)
+        self.act = nn.PReLU() if act == 'prelu' else act
+
+        if bias:
+            self.bias = nn.Parameter(torch.FloatTensor(out_ft))
+            self.bias.data.fill_(0.0)
+        else:
+            self.register_parameter('bias', None)
+
+        for m in self.modules():
+            self.weights_init(m)
+
+    def weights_init(self, m):
+        if isinstance(m, nn.Linear):
+            torch.nn.init.xavier_uniform_(m.weight.data)
+            if m.bias is not None:
+                m.bias.data.fill_(0.0)
+
+    def forward(self, seq, adj, sparse=False):
+        seq_fts = self.fc(seq)
+        if sparse:
+            out = torch.unsqueeze(torch.spmm(adj, torch.squeeze(seq_fts, 0)),
+                                  0)
+        else:
+            out = torch.bmm(adj, seq_fts)
+        if self.bias is not None:
+            out += self.bias
+
+        return self.act(out)
+
+
+class AvgReadout(nn.Module):
+    def __init__(self):
+        super(AvgReadout, self).__init__()
+
+    def forward(self, seq):
+        return torch.mean(seq, 1)
+
+
+class MaxReadout(nn.Module):
+    def __init__(self):
+        super(MaxReadout, self).__init__()
+
+    def forward(self, seq):
+        return torch.max(seq, 1).values
+
+
+class MinReadout(nn.Module):
+    def __init__(self):
+        super(MinReadout, self).__init__()
+
+    def forward(self, seq):
+        return torch.min(seq, 1).values
+
+
+class WSReadout(nn.Module):
+    def __init__(self):
+        super(WSReadout, self).__init__()
+
+    def forward(self, seq, query):
+        query = query.permute(0, 2, 1)
+        sim = torch.matmul(seq, query)
+        sim = F.softmax(sim, dim=1)
+        sim = sim.repeat(1, 1, 64)
+        out = torch.mul(seq, sim)
+        out = torch.sum(out, 1)
+        return out
+
+
+class Discriminator(nn.Module):
+    def __init__(self, n_h, negsamp_round):
+        super(Discriminator, self).__init__()
+        self.f_k = nn.Bilinear(n_h, n_h, 1)
+
+        for m in self.modules():
+            self.weights_init(m)
+
+        self.negsamp_round = negsamp_round
+
+    def weights_init(self, m):
+        if isinstance(m, nn.Bilinear):
+            torch.nn.init.xavier_uniform_(m.weight.data)
+            if m.bias is not None:
+                m.bias.data.fill_(0.0)
+
+    def forward(self, c, h_pl):
+        scs = []
+        # positive
+        scs.append(self.f_k(h_pl, c))
+
+        # negative
+        c_mi = c
+        for _ in range(self.negsamp_round):
+            c_mi = torch.cat((c_mi[-2:-1, :], c_mi[:-1, :]), 0)
+            scs.append(self.f_k(h_pl, c_mi))
+
+        logits = torch.cat(tuple(scs))
+
+        return logits
+
+
+class Contextual_Discriminator(nn.Module):
+    def __init__(self, n_h, negsamp_round):
+        super(Contextual_Discriminator, self).__init__()
+        self.f_k = nn.Bilinear(n_h, n_h, 1)
+        for m in self.modules():
+            self.weights_init(m)
+        self.negsamp_round = negsamp_round
+
+    def weights_init(self, m):
+        if isinstance(m, nn.Bilinear):
+            torch.nn.init.xavier_uniform_(m.weight.data)
+            if m.bias is not None:
+                m.bias.data.fill_(0.0)
+
+    def forward(self, c, h_pl, s_bias1=None, s_bias2=None):
+        scs = []
+        scs.append(self.f_k(h_pl, c))
+        c_mi = c
+        for _ in range(self.negsamp_round):
+            c_mi = torch.cat((c_mi[-2:-1, :], c_mi[:-1, :]), 0)
+            scs.append(self.f_k(h_pl, c_mi))
+        logits = torch.cat(tuple(scs))
+        return logits
+
+
+class Patch_Discriminator(nn.Module):
+    def __init__(self, n_h, negsamp_round):
+        super(Patch_Discriminator, self).__init__()
+        self.f_k = nn.Bilinear(n_h, n_h, 1)
+        for m in self.modules():
+            self.weights_init(m)
+        self.negsamp_round = negsamp_round
+
+    def weights_init(self, m):
+        if isinstance(m, nn.Bilinear):
+            torch.nn.init.xavier_uniform_(m.weight.data)
+            if m.bias is not None:
+                m.bias.data.fill_(0.0)
+
+    def forward(self, h_ano, h_unano, s_bias1=None, s_bias2=None):
+        scs = []
+        scs.append(self.f_k(h_unano, h_ano))
+        h_mi = h_ano
+        for _ in range(self.negsamp_round):
+            h_mi = torch.cat((h_mi[-2:-1, :], h_mi[:-1, :]), 0)
+            scs.append(self.f_k(h_unano, h_mi))
+        logits = torch.cat(tuple(scs))
+        return logits
diff --git a/pygod/models/anomalydae.py b/pygod/models/anomalydae.py
index 14ed8cd..a6f240c 100644
--- a/pygod/models/anomalydae.py
+++ b/pygod/models/anomalydae.py
@@ -5,174 +5,19 @@
 
 
 import torch
+import numpy as np
 import torch.nn as nn
 import torch.nn.functional as F
-from torch_sparse import SparseTensor
 from torch_geometric.nn import GATConv
+from torch_geometric.utils import to_dense_adj
+from torch_geometric.loader import NeighborLoader
 from sklearn.utils.validation import check_is_fitted
 
 from . import BaseDetector
+from ..utils.utility import validate_device
 from ..utils.metric import eval_roc_auc
 
 
-class StructureAE(nn.Module):
-    """
-    Description
-    -----------
-    Structure Autoencoder in AnomalyDAE model: the encoder
-    transforms the node attribute X into the latent
-    representation with the linear layer, and a graph attention
-    layer produces an embedding with weight importance of node 
-    neighbors. Finally, the decoder reconstructs the final embedding
-    to the original.
-
-    See :cite:`fan2020anomalydae` for details.
-
-    Parameters
-    ----------
-    in_dim: int
-        input dimension of node data
-    embed_dim: int
-        the latent representation dimension of node
-       (after the first linear layer)
-    out_dim: int 
-        the output dim after the graph attention layer
-    dropout: float
-        dropout probability for the linear layer
-    act: F, optional
-         Choice of activation function        
-    Returns
-    -------
-    x : torch.Tensor
-        Reconstructed attribute (feature) of nodes.
-    embed_x : torch.Tensor
-              Embedd nodes after the attention layer
-    """
-
-    def __init__(self,
-                 in_dim,
-                 embed_dim,
-                 out_dim,
-                 dropout,
-                 act):
-        super(StructureAE, self).__init__()
-        self.dense = nn.Linear(in_dim, embed_dim)
-        self.attention_layer = GATConv(embed_dim, out_dim)
-        self.dropout = dropout
-        self.act = act
-
-    def forward(self,
-                x,
-                edge_index):
-        # encoder
-        x = self.act(self.dense(x))
-        x = F.dropout(x, self.dropout)
-        # print(x.shape, adj.shape)
-        embed_x = self.attention_layer(x, edge_index)
-        # decoder
-        x = torch.sigmoid(embed_x @ embed_x.T)
-        return x, embed_x
-
-
-class AttributeAE(nn.Module):
-    """
-    Description
-    -----------
-    Attribute Autoencoder in AnomalyDAE model: the encoder
-    employs two non-linear feature transform to the node attribute
-    x. The decoder takes both the node embeddings from the structure
-    autoencoder and the reduced attribute representation to 
-    reconstruct the original node attribute.
-
-    Parameters
-    ----------
-    in_dim:  int
-        dimension of the input number of nodes
-    embed_dim: int
-        the latent representation dimension of node
-        (after the first linear layer)
-    out_dim:  int
-        the output dim after two linear layers
-    dropout: float
-        dropout probability for the linear layer
-    act: F, optional
-         Choice of activation function   
-    Returns
-    -------
-    x : torch.Tensor
-        Reconstructed attribute (feature) of nodes.
-    """
-
-    def __init__(self,
-                 in_dim,
-                 embed_dim,
-                 out_dim,
-                 dropout,
-                 act):
-        super(AttributeAE, self).__init__()
-        self.dense1 = nn.Linear(in_dim, embed_dim)
-        self.dense2 = nn.Linear(embed_dim, out_dim)
-        self.dropout = dropout
-        self.act = act
-
-    def forward(self,
-                x,
-                struct_embed):
-        # encoder
-        x = self.act(self.dense1(x.T))
-        x = F.dropout(x, self.dropout)
-        x = self.dense2(x)
-        x = F.dropout(x, self.dropout)
-        # decoder
-        x = struct_embed @ x.T
-        return x
-
-
-class AnomalyDAE_Base(nn.Module):
-    """
-    Description
-    -----------
-    AdnomalyDAE_Base is an anomaly detector consisting of a structure autoencoder,
-    and an attribute reconstruction autoencoder. 
-
-    Parameters
-    ----------
-    in_node_dim : int
-         Dimension of input feature
-    in_num_dim: int
-         Dimension of the input number of nodes
-    embed_dim:: int
-         Dimension of the embedding after the first reduced linear layer (D1)   
-    out_dim : int
-         Dimension of final representation
-    dropout : float, optional
-        Dropout rate of the model
-        Default: 0
-    act: F, optional
-         Choice of activation function
-    """
-
-    def __init__(self,
-                 in_node_dim,
-                 in_num_dim,
-                 embed_dim,
-                 out_dim,
-                 dropout,
-                 act):
-        super(AnomalyDAE_Base, self).__init__()
-        self.structure_AE = StructureAE(in_node_dim, embed_dim,
-                                        out_dim, dropout, act)
-        self.attribute_AE = AttributeAE(in_num_dim, embed_dim,
-                                        out_dim, dropout, act)
-
-    def forward(self,
-                x,
-                edge_index):
-        A_hat, embed_x = self.structure_AE(x, edge_index)
-        X_hat = self.attribute_AE(x, embed_x)
-        return A_hat, X_hat
-
-
 class AnomalyDAE(BaseDetector):
     """
     AnomalyDAE (Dual autoencoder for anomaly detection on attributed networks):
@@ -221,6 +66,11 @@ class AnomalyDAE(BaseDetector):
         Maximum number of training epoch. Defaults: ``5``.
     gpu : int
         GPU Index, -1 for using CPU. Defaults: ``0``.
+    batch_size : int, optional
+        Minibatch size, 0 for full batch training. Default: ``0``.
+    num_neigh : int, optional
+        Number of neighbors in sampling, -1 for all neighbors.
+        Default: ``-1``.
     verbose : bool
         Verbosity mode. Turn on to print out log information.
         Defaults: ``False``.
@@ -246,6 +96,8 @@ def __init__(self,
                  lr=0.004,
                  epoch=5,
                  gpu=0,
+                 batch_size=0,
+                 num_neigh=-1,
                  verbose=False):
         super(AnomalyDAE, self).__init__(contamination=contamination)
 
@@ -262,10 +114,9 @@ def __init__(self,
         # training param
         self.lr = lr
         self.epoch = epoch
-        if gpu >= 0 and torch.cuda.is_available():
-            self.device = 'cuda:{}'.format(gpu)
-        else:
-            self.device = 'cpu'
+        self.device = validate_device(gpu)
+        self.batch_size = batch_size
+        self.num_neigh = num_neigh
 
         # other param
         self.verbose = verbose
@@ -273,27 +124,32 @@ def __init__(self,
 
     def fit(self, G, y_true=None):
         """
-        Description
-        -----------
         Fit detector with input data.
 
         Parameters
         ----------
-        G : PyTorch Geometric Data instance (torch_geometric.data.Data)
+        G : torch_geometric.data.Data
             The input data.
-        y_true : numpy.array, optional (default=None)
-            The optional outlier ground truth labels used to monitor the
-            training progress. They are not used to optimize the
-            unsupervised model.
+        y_true : numpy.ndarray, optional
+            The optional outlier ground truth labels used to monitor
+            the training progress. They are not used to optimize the
+            unsupervised model. Default: ``None``.
 
         Returns
         -------
         self : object
             Fitted estimator.
         """
-        attrs, adj, edge_index = self.process_graph(G)
-        self.model = AnomalyDAE_Base(in_node_dim=attrs.shape[1],
-                                     in_num_dim=attrs.shape[0],
+        G.node_idx = torch.arange(G.x.shape[0])
+        G.s = to_dense_adj(G.edge_index)[0]
+        if self.batch_size == 0:
+            self.batch_size = G.x.shape[0]
+        loader = NeighborLoader(G,
+                                [self.num_neigh],
+                                batch_size=self.batch_size)
+
+        self.model = AnomalyDAE_Base(in_node_dim=G.x.shape[1],
+                                     in_num_dim=self.batch_size,
                                      embed_dim=self.embed_dim,
                                      out_dim=self.out_dim,
                                      dropout=self.dropout,
@@ -303,38 +159,43 @@ def fit(self, G, y_true=None):
                                      lr=self.lr,
                                      weight_decay=self.weight_decay)
 
+        self.model.train()
+        decision_scores = np.zeros(G.x.shape[0])
         for epoch in range(self.epoch):
-            self.model.train()
-            optimizer.zero_grad()
-            A_hat, X_hat = self.model(attrs, edge_index)
-            loss, struct_loss, feat_loss = self.loss_func(adj, A_hat,
-                                                          attrs, X_hat)
-            loss_mean = torch.mean(loss)
-            loss_mean.backward()
-            optimizer.step()
+            epoch_loss = 0
+            for sampled_data in loader:
+                batch_size = sampled_data.batch_size
+                node_idx = sampled_data.node_idx
+                x, s, edge_index = self.process_graph(sampled_data)
+
+                x_, s_ = self.model(x, edge_index, batch_size)
+                score = self.loss_func(x[:batch_size],
+                                       x_[:batch_size],
+                                       s[:batch_size, node_idx],
+                                       s_[:batch_size])
+                decision_scores[node_idx[:batch_size]] = score.detach() \
+                    .cpu().numpy()
+                loss = torch.mean(score)
+                epoch_loss += loss.item() * batch_size
+
+                optimizer.zero_grad()
+                loss.backward()
+                optimizer.step()
 
             if self.verbose:
-                print("Epoch:", '%04d' % epoch, "train_loss=",
-                      "{:.5f}".format(loss_mean.item()), "train/struct_loss=",
-                      "{:.5f}".format(struct_loss.item()), "train/feat_loss=",
-                      "{:.5f}".format(feat_loss.item()))
-
-            self.model.eval()
-            A_hat, X_hat = self.model(attrs, edge_index)
-            loss, struct_loss, feat_loss = self.loss_func(adj, A_hat,
-                                                          attrs, X_hat)
-            score = loss.detach().cpu().numpy()
-            if self.verbose:
-                print('AUC', eval_roc_auc(y_true, score))
-
-        self.decision_scores_ = score
+                print("Epoch {:04d}: Loss {:.4f}"
+                      .format(epoch, epoch_loss / G.x.shape[0]), end='')
+                if y_true is not None:
+                    auc = eval_roc_auc(y_true, decision_scores)
+                    print(" | AUC {:.4f}".format(auc), end='')
+                print()
+
+        self.decision_scores_ = decision_scores
         self._process_decision_scores()
         return self
 
     def decision_function(self, G):
         """
-        Description
-        -----------
         Predict raw anomaly score of X using the fitted detector.
         The anomaly score of an input sample is computed based on different
         detector algorithms. For consistency, outliers are assigned with
@@ -350,24 +211,34 @@ def decision_function(self, G):
         anomaly_scores : numpy.ndarray
             The anomaly score of shape :math:`N`.
         """
-
         check_is_fitted(self, ['model'])
+        G.node_idx = torch.arange(G.x.shape[0])
+        G.s = to_dense_adj(G.edge_index)[0]
 
-        # get needed data object from the input data
-        attrs, adj, edge_index = self.process_graph(G)
+        loader = NeighborLoader(G,
+                                [self.num_neigh],
+                                batch_size=self.batch_size)
 
-        # enable the evaluation mode
         self.model.eval()
+        outlier_scores = np.zeros(G.x.shape[0])
+        for sampled_data in loader:
+            batch_size = sampled_data.batch_size
+            node_idx = sampled_data.node_idx
+
+            x, s, edge_index = self.process_graph(sampled_data)
+
+            x_, s_ = self.model(x, edge_index, batch_size)
+            score = self.loss_func(x[:batch_size],
+                                   x_[:batch_size],
+                                   s[:batch_size, node_idx],
+                                   s_[:batch_size])
 
-        # construct the vector for holding the reconstruction error
-        A_hat, X_hat = self.model(attrs, edge_index)
-        outlier_scores, _, _ = self.loss_func(adj, A_hat, attrs, X_hat)
-        return outlier_scores.detach().cpu().numpy()
+            outlier_scores[node_idx[:batch_size]] = score.detach() \
+                .cpu().numpy()
+        return outlier_scores
 
     def process_graph(self, G):
         """
-        Description
-        -----------
         Process the raw PyG data object into a tuple of sub data
         objects needed for the model.
 
@@ -380,60 +251,193 @@ def process_graph(self, G):
         -------
         x : torch.Tensor
             Attribute (feature) of nodes.
-        adj : torch.Tensor
+        s : torch.Tensor
             Adjacency matrix of the graph.
         edge_index : torch.Tensor
             Edge list of the graph.
         """
-        edge_index = G.edge_index
-
-        #  via sparse matrix operation
-        dense_adj \
-            = SparseTensor(row=edge_index[0], col=edge_index[1]).to_dense()
-
-        # adjacency matrix normalization
-        rowsum = dense_adj.sum(1)
-        d_inv_sqrt = torch.pow(rowsum, -0.5).flatten()
-        d_inv_sqrt[torch.isinf(d_inv_sqrt)] = 0.
-        d_mat_inv_sqrt = torch.diag(d_inv_sqrt)
-        adj = (dense_adj * d_mat_inv_sqrt).T * d_mat_inv_sqrt
-
-        edge_index = edge_index.to(self.device)
-        adj = adj.to(self.device)
+        s = G.s.to(self.device)
+        edge_index = G.edge_index.to(self.device)
         x = G.x.to(self.device)
 
-        # return data objects needed for the network
-        return x, adj, edge_index
+        return x, s, edge_index
 
-    def loss_func(self,
-                  adj,
-                  A_hat,
-                  attrs,
-                  X_hat):
+    def loss_func(self, x, x_, s, s_):
         # generate hyperparameter - structure penalty
-        reversed_adj = torch.ones(adj.shape).to(self.device) - adj
+        reversed_adj = 1 - s
         thetas = torch.where(reversed_adj > 0, reversed_adj,
-                             torch.full(adj.shape, self.theta).to(self.device))
+                             torch.full(s.shape, self.theta).to(self.device))
 
         # generate hyperparameter - node penalty
-        reversed_attr = torch.ones(attrs.shape).to(self.device) - attrs
+        reversed_attr = 1 - x
         etas = torch.where(reversed_attr == 1, reversed_attr,
-                           torch.full(attrs.shape, self.eta).to(self.device))
+                           torch.full(x.shape, self.eta).to(self.device))
 
-        # Attribute reconstruction loss
-        diff_attribute = torch.pow(X_hat -
-                                   attrs, 2) * etas
-        attribute_reconstruction_errors = \
-            torch.sqrt(torch.sum(diff_attribute, 1))
-        attribute_cost = torch.mean(attribute_reconstruction_errors)
+        # attribute reconstruction loss
+        diff_attribute = torch.pow(x_ - x, 2) * etas
+        attribute_errors = torch.sqrt(torch.sum(diff_attribute, 1))
 
         # structure reconstruction loss
-        diff_structure = torch.pow(A_hat - adj, 2) * thetas
-        structure_reconstruction_errors = \
-            torch.sqrt(torch.sum(diff_structure, 1))
-        structure_cost = torch.mean(structure_reconstruction_errors)
+        diff_structure = torch.pow(s_ - s, 2) * thetas
+        structure_errors = torch.sqrt(torch.sum(diff_structure, 1))
+
+        score = self.alpha * attribute_errors \
+                + (1 - self.alpha) * structure_errors
+
+        return score
+
+
+class AnomalyDAE_Base(nn.Module):
+    """
+    AnomalyDAE_Base is an anomaly detector consisting of a structure
+    autoencoder and an attribute reconstruction autoencoder.
+
+    Parameters
+    ----------
+    in_node_dim : int
+         Dimension of input feature
+    in_num_dim: int
+         Dimension of the input number of nodes
+    embed_dim:: int
+         Dimension of the embedding after the first reduced linear
+         layer (D1)
+    out_dim : int
+         Dimension of final representation
+    dropout : float, optional
+        Dropout rate of the model
+        Default: 0
+    act: F, optional
+         Choice of activation function
+    """
+
+    def __init__(self,
+                 in_node_dim,
+                 in_num_dim,
+                 embed_dim,
+                 out_dim,
+                 dropout,
+                 act):
+        super(AnomalyDAE_Base, self).__init__()
+
+        self.num_center_nodes = in_num_dim
+        self.structure_ae = StructureAE(in_node_dim,
+                                        embed_dim,
+                                        out_dim,
+                                        dropout,
+                                        act)
+        self.attribute_ae = AttributeAE(self.num_center_nodes,
+                                        embed_dim,
+                                        out_dim,
+                                        dropout,
+                                        act)
+
+    def forward(self, x, edge_index, batch_size):
+        s_, h = self.structure_ae(x, edge_index)
+        if batch_size < self.num_center_nodes:
+            x = F.pad(x, (0, 0, 0, self.num_center_nodes - batch_size))
+        x_ = self.attribute_ae(x[:self.num_center_nodes], h)
+        return x_, s_
+
+
+class StructureAE(nn.Module):
+    """
+    Structure Autoencoder in AnomalyDAE model: the encoder
+    transforms the node attribute X into the latent
+    representation with the linear layer, and a graph attention
+    layer produces an embedding with weight importance of node
+    neighbors. Finally, the decoder reconstructs the final embedding
+    to the original.
+
+    Parameters
+    ----------
+    in_dim: int
+        input dimension of node data
+    embed_dim: int
+        the latent representation dimension of node
+       (after the first linear layer)
+    out_dim: int
+        the output dim after the graph attention layer
+    dropout: float
+        dropout probability for the linear layer
+    act: F, optional
+         Choice of activation function
 
-        cost = self.alpha * attribute_reconstruction_errors + (
-                1 - self.alpha) * structure_reconstruction_errors
+    Returns
+    -------
+    x : torch.Tensor
+        Reconstructed attribute (feature) of nodes.
+    embed_x : torch.Tensor
+        Embed nodes after the attention layer
+    """
 
-        return cost, structure_cost, attribute_cost
+    def __init__(self,
+                 in_dim,
+                 embed_dim,
+                 out_dim,
+                 dropout,
+                 act):
+        super(StructureAE, self).__init__()
+        self.dense = nn.Linear(in_dim, embed_dim)
+        self.attention_layer = GATConv(embed_dim, out_dim)
+        self.dropout = dropout
+        self.act = act
+
+    def forward(self, x, edge_index):
+        # encoder
+        x = self.act(self.dense(x))
+        x = F.dropout(x, self.dropout)
+        h = self.attention_layer(x, edge_index)
+        # decoder
+        s_ = torch.sigmoid(h @ h.T)
+        return s_, h
+
+
+class AttributeAE(nn.Module):
+    """
+    Attribute Autoencoder in AnomalyDAE model: the encoder
+    employs two non-linear feature transform to the node attribute
+    x. The decoder takes both the node embeddings from the structure
+    autoencoder and the reduced attribute representation to
+    reconstruct the original node attribute.
+
+    Parameters
+    ----------
+    in_dim:  int
+        dimension of the input number of nodes
+    embed_dim: int
+        the latent representation dimension of node
+        (after the first linear layer)
+    out_dim:  int
+        the output dim after two linear layers
+    dropout: float
+        dropout probability for the linear layer
+    act: F, optional
+         Choice of activation function
+
+    Returns
+    -------
+    x : torch.Tensor
+        Reconstructed attribute (feature) of nodes.
+    """
+
+    def __init__(self,
+                 in_dim,
+                 embed_dim,
+                 out_dim,
+                 dropout,
+                 act):
+        super(AttributeAE, self).__init__()
+        self.dense1 = nn.Linear(in_dim, embed_dim)
+        self.dense2 = nn.Linear(embed_dim, out_dim)
+        self.dropout = dropout
+        self.act = act
+
+    def forward(self, x, h):
+        # encoder
+        x = self.act(self.dense1(x.T))
+        x = F.dropout(x, self.dropout)
+        x = self.dense2(x)
+        x = F.dropout(x, self.dropout)
+        # decoder
+        x = h @ x.T
+        return x
diff --git a/pygod/models/base.py b/pygod/models/base.py
index 8bf30af..c4e17cd 100644
--- a/pygod/models/base.py
+++ b/pygod/models/base.py
@@ -322,21 +322,21 @@ def _get_param_names(cls):
 
     # noinspection PyPep8
     def get_params(self, deep=True):
-        """Get parameters for this estimator.
-        See https://scikit-learn.org/stable/modules/generated/sklearn.base
-        .BaseEstimator.html
-        and sklearn/base.py for more information.
+        r"""Get parameters for this estimator. See `sklearn.base.BaseEstimator
+        <https://scikit-learn.org/stable/modules/generated/sklearn.base.BaseEstimator.html>`_
+        for more information.
+
         Parameters
         ----------
-        deep : bool, optional (default=True)
+        deep : bool, optional
             If True, will return the parameters for this estimator and
             contained sub-objects that are estimators.
+            Default: ```True```.
 
         Returns
         -------
         params : mapping of string to any
             Parameter names mapped to their values.
-
         """
 
         out = dict()
@@ -369,9 +369,10 @@ def set_params(self, **params):
         (such as pipelines). The latter have parameters of the form
         ``<component>__<parameter>`` so that it's possible to update each
         component of a nested object.
-        See https://scikit-learn.org/stable/modules/generated/sklearn.base
-        .BaseEstimator.html
-        and sklearn/base.py for more information.
+        See `sklearn.base.BaseEstimator
+        <https://scikit-learn.org/stable/modules/generated/sklearn.base.BaseEstimator.html>`_
+        for more information.
+
         Returns
         -------
         self : object
@@ -403,12 +404,6 @@ def set_params(self, **params):
 
     def __repr__(self):
         # noinspection PyPep8
-        """
-        See https://scikit-learn.org/stable/modules/generated/sklearn.base
-        .BaseEstimator.html
-        and sklearn/base.py for more information.
-        """
-
         class_name = self.__class__.__name__
         return '%s(%s)' % (class_name, _pprint(self.get_params(deep=False),
                                                offset=len(class_name), ),)
diff --git a/pygod/models/basic_nn.py b/pygod/models/basic_nn.py
index 9971711..1b048ea 100644
--- a/pygod/models/basic_nn.py
+++ b/pygod/models/basic_nn.py
@@ -275,3 +275,37 @@ def forward(self, x: Tensor, edge_index: Adj, *args, **kwargs) -> Tensor:
     def __repr__(self) -> str:
         return (f'{self.__class__.__name__}({self.in_channels}, '
                 f'{self.out_channels}, num_layers={self.num_layers})')
+
+
+class Vanilla_GCN(torch.nn.Module):
+    def __init__(self, in_ft, out_ft, act, bias=True):
+        super(Vanilla_GCN, self).__init__()
+        self.fc = torch.nn.Linear(in_ft, out_ft, bias=False)
+        self.act = torch.nn.PReLU() if act == 'prelu' else act
+
+        if bias:
+            self.bias = torch.nn.Parameter(torch.FloatTensor(out_ft))
+            self.bias.data.fill_(0.0)
+        else:
+            self.register_parameter('bias', None)
+
+        for m in self.modules():
+            self.weights_init(m)
+
+    def weights_init(self, m):
+        if isinstance(m, torch.nn.Linear):
+            torch.nn.init.xavier_uniform_(m.weight.data)
+            if m.bias is not None:
+                m.bias.data.fill_(0.0)
+
+    def forward(self, seq, adj, sparse=False):
+        seq_fts = self.fc(seq)
+        if sparse:
+            out = torch.unsqueeze(torch.spmm(adj, torch.squeeze(seq_fts, 0)),
+                                  0)
+        else:
+            out = torch.bmm(adj, seq_fts)
+        if self.bias is not None:
+            out += self.bias
+
+        return self.act(out)
diff --git a/pygod/models/cola.py b/pygod/models/cola.py
new file mode 100644
index 0000000..c0cdeef
--- /dev/null
+++ b/pygod/models/cola.py
@@ -0,0 +1,419 @@
+# -*- coding: utf-8 -*-
+"""Anomaly Detection on Attributed Networks via Contrastive
+Self-Supervised Learning (CoLA)
+The code is adapted from the source code from paper authors:
+https://github.com/GRAND-Lab/CoLA"""
+# Author: Canyu Chen <cchen151@hawk.iit.edu>
+# License: BSD 2 clause
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+import scipy.sparse as sp
+import random
+import os
+
+from torch_geometric.utils import to_dense_adj
+from torch_cluster import random_walk
+
+from . import BaseDetector
+from .basic_nn import Vanilla_GCN as GCN
+from ..utils.utility import validate_device
+from ..utils.metric import eval_roc_auc
+
+os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"
+
+
+def generate_rw_subgraph(pyg_graph, nb_nodes, subgraph_size):
+    """Generate subgraph with random walk algorithm."""
+    row, col = pyg_graph.edge_index
+    all_idx = torch.tensor(list(range(nb_nodes)))
+    traces = random_walk(row, col, all_idx, walk_length=3)
+    subv = traces.tolist()
+    return subv
+
+
+class CoLA(BaseDetector):
+    """
+    CoLA (Anomaly Detection on Attributed Networks via
+    Contrastive Self-Supervised Learning)
+    CoLA is a contrastive self-supervised learning-based method
+    for graph anomaly detection. (beta)
+
+    Parameters
+    ----------
+    lr : float, optional
+        Learning rate. Default: ``1e-3``.
+    epoch : int, optional
+        Maximum number of training epoch. Default: ``10``.
+    embedding_dim : int, optional
+        The node embedding dimension obtained by the GCN module of CoLA.
+        Default: ``64``.
+    negsamp_ratio : int, optional
+        Number of negative samples for each instance used by the contrastive
+        learning module. Default: ``1``.
+    readout : str, optional
+        The readout layer type used by CoLA model. Default: ``avg`` .
+    weight_decay : float, optional
+        Weight decay (L2 penalty). Default: ``0.``.
+    batch_size : int, optional
+        Minibatch size, 0 for full batch training. Default: ``0``.
+    subgraph_size : int, optional
+        Number of nodes in the subgraph sampled by random walk. Default: ``4``.
+    contamination : float, optional
+        Valid in (0., 0.5). The proportion of outliers in the data set.
+        Used when fitting to define the threshold on the decision
+        function. Default: ``0.1``.
+    gpu : int, optional
+        GPU Index, -1 for using CPU. Default: ``0``.
+    verbose : bool, optional
+        Verbosity mode. Turn on to print out log information.
+        Default: ``False``.
+
+    Examples
+    --------
+    >>> from pygod.models import CoLA
+    >>> model = CoLA()
+    >>> model.fit(data)
+    >>> prediction = model.predict(data)
+    """
+
+    def __init__(self,
+                 lr=1e-3,
+                 epoch=10,
+                 embedding_dim=64,
+                 negsamp_ratio=1,
+                 readout='avg',
+                 weight_decay=0.,
+                 batch_size=0,
+                 subgraph_size=4,
+                 contamination=0.1,
+                 gpu=0,
+                 verbose=False):
+        super(CoLA, self).__init__(contamination=contamination)
+
+        self.lr = lr
+        self.num_epoch = epoch
+        self.embedding_dim = embedding_dim
+        self.negsamp_ratio = negsamp_ratio
+        self.readout = readout
+        self.weight_decay = weight_decay
+        self.batch_size = batch_size
+        self.subgraph_size = subgraph_size
+        self.device = validate_device(gpu)
+
+        self.verbose = verbose
+        self.model = None
+
+    def fit(self, G, y_true=None):
+        """
+        Fit detector with input data.
+
+        Parameters
+        ----------
+        G : PyTorch Geometric Data instance (torch_geometric.data.Data)
+            The input data.
+        y_true : numpy.ndarray, optional
+            The optional outlier ground truth labels used to monitor
+            the training progress. They are not used to optimize the
+            unsupervised model. Default: ``None``.
+
+        Returns
+        -------
+        self : object
+            Fitted estimator.
+        """
+
+        x, adj = self.process_graph(G)
+
+        # Initialize model and optimiser
+        self.model = CoLA_Base(self.feat_dim,
+                               self.embedding_dim,
+                               'prelu',
+                               self.negsamp_ratio,
+                               self.readout,
+                               self.subgraph_size,
+                               self.device).to(self.device)
+
+        optimizer = torch.optim.Adam(self.model.parameters(),
+                                     lr=self.lr,
+                                     weight_decay=self.weight_decay)
+
+        if self.batch_size:
+            batch_num = self.num_nodes // self.batch_size + 1
+        else:  # full batch training
+            batch_num = 1
+
+        multi_epoch_ano_score = np.zeros((self.num_epoch, self.num_nodes))
+
+        for epoch in range(self.num_epoch):
+
+            self.model.train()
+            epoch_loss = 0
+            all_idx = list(range(self.num_nodes))
+            random.shuffle(all_idx)
+            decision_scores = np.zeros(self.num_nodes)
+
+            subgraphs = generate_rw_subgraph(G, self.num_nodes,
+                                             self.subgraph_size)
+
+            for batch_idx in range(batch_num):
+
+                is_final_batch = (batch_idx == (batch_num - 1))
+
+                if not is_final_batch:
+                    idx = all_idx[
+                          batch_idx *
+                          self.batch_size: (batch_idx + 1) * self.batch_size]
+                else:
+                    idx = all_idx[batch_idx * self.batch_size:]
+
+                cur_batch_size = len(idx)
+
+                output = self.model(x, adj, idx, subgraphs, cur_batch_size)
+                loss = self.loss_function(output, cur_batch_size)
+
+                loss = torch.mean(loss)
+                epoch_loss += loss.item() * cur_batch_size
+
+                optimizer.zero_grad()
+                loss.backward()
+                optimizer.step()
+
+                # compute the anomaly score
+                logits = torch.sigmoid(torch.squeeze(output))
+                ano_score = - (logits[:cur_batch_size] - logits[
+                                cur_batch_size:]).detach().cpu().numpy()
+                decision_scores[idx] = ano_score
+
+            multi_epoch_ano_score[epoch, :] = decision_scores
+
+            if self.verbose:
+                print("Epoch {:04d}: Loss {:.4f}"
+                      .format(epoch, epoch_loss / G.x.shape[0]), end='')
+                if y_true is not None:
+                    auc = eval_roc_auc(y_true, decision_scores)
+                    print(" | AUC {:.4f}".format(auc), end='')
+                print()
+
+        ano_score_final = np.mean(multi_epoch_ano_score, axis=0)
+
+        self.decision_scores_ = ano_score_final
+        self._process_decision_scores()
+        return self
+
+    def decision_function(self, G, rounds=10):
+        """
+        Predict raw anomaly score using the fitted detector. Outliers
+        are assigned with larger anomaly scores.
+
+        Parameters
+        ----------
+        G : PyTorch Geometric Data instance (torch_geometric.data.Data)
+            The input data.
+        rounds : int, optional
+            Number of rounds to generate the decision score. Default: ``10``.
+
+        Returns
+        -------
+        outlier_scores : numpy.ndarray
+            The anomaly score of shape :math:`N`.
+        """
+
+        x, adj = self.process_graph(G)
+
+        if self.batch_size:
+            batch_num = self.num_nodes // self.batch_size + 1
+        else:  # full batch training
+            batch_num = 1
+
+        multi_round_ano_score = np.zeros((rounds, self.num_nodes))
+
+        # enable the evaluation mode
+        self.model.eval()
+
+        for r in range(rounds):
+
+            all_idx = list(range(self.num_nodes))
+            random.shuffle(all_idx)
+
+            subgraphs = generate_rw_subgraph(G, self.num_nodes,
+                                             self.subgraph_size)
+
+            for batch_idx in range(batch_num):
+
+                is_final_batch = (batch_idx == (batch_num - 1))
+
+                if not is_final_batch:
+                    idx = all_idx[
+                          batch_idx *
+                          self.batch_size: (batch_idx + 1) * self.batch_size]
+                else:
+                    idx = all_idx[batch_idx * self.batch_size:]
+
+                cur_batch_size = len(idx)
+
+                with torch.no_grad():
+                    output = self.model(x, adj, idx, subgraphs, cur_batch_size)
+                logits = torch.sigmoid(torch.squeeze(output))
+                ano_score = - (logits[:cur_batch_size] - logits[
+                                cur_batch_size:]).cpu().numpy()
+                multi_round_ano_score[r, idx] = ano_score
+
+        ano_score_final = np.mean(multi_round_ano_score, axis=0)
+
+        return ano_score_final
+
+    def process_graph(self, G):
+        """
+        Description
+        -----------
+        Process the raw PyG data object into a tuple of sub data
+        objects needed for the model.
+
+        Parameters
+        ----------
+        G : PyTorch Geometric Data instance (torch_geometric.data.Data)
+            The input data.
+
+        Returns
+        -------
+        x : torch.Tensor
+            Attribute (feature) of nodes.
+        adj : torch.Tensor
+            Adjacency matrix of the graph.
+        """
+        self.num_nodes = G.x.shape[0]
+        self.feat_dim = G.x.shape[1]
+        adj = to_dense_adj(G.edge_index)[0]
+        adj = sp.coo_matrix(adj)
+        rowsum = np.array(adj.sum(1))
+        d_inv_sqrt = np.power(rowsum, -0.5).flatten()
+        d_inv_sqrt[np.isinf(d_inv_sqrt)] = 0.
+        d_mat_inv_sqrt = sp.diags(d_inv_sqrt)
+        adj = adj.dot(d_mat_inv_sqrt).transpose().dot(d_mat_inv_sqrt).tocoo()
+        adj = (adj + sp.eye(adj.shape[0])).todense()
+
+        x = torch.FloatTensor(G.x[np.newaxis]).to(self.device)
+        adj = torch.FloatTensor(adj[np.newaxis]).to(self.device)
+
+        # return data objects needed for the network
+        return x, adj
+
+    def loss_function(self, logits, batch_size):
+
+        b_xent = nn.BCEWithLogitsLoss(reduction='none',
+                                      pos_weight=torch.tensor(
+                                          [self.negsamp_ratio]))
+
+        lbl = torch.unsqueeze(
+            torch.cat((torch.ones(batch_size),
+                       torch.zeros(batch_size * self.negsamp_ratio))), 1)
+
+        score = b_xent(logits.cpu(), lbl.cpu())
+
+        return score
+
+
+class CoLA_Base(nn.Module):
+    def __init__(self,
+                 n_in,
+                 n_h,
+                 activation,
+                 negsamp_round,
+                 readout,
+                 subgraph_size,
+                 device):
+
+        super(CoLA_Base, self).__init__()
+        self.n_in = n_in
+        self.subgraph_size = subgraph_size
+        self.device = device
+        self.readout = readout
+        self.gcn = GCN(n_in, n_h, activation).to(self.device)
+        self.disc = Discriminator(n_h, negsamp_round)
+
+    def forward(self, x, adj, idx, subgraphs, batch_size, sparse=False):
+
+        batch_adj = []
+        batch_feature = []
+        added_adj_zero_row = torch.zeros(
+            (batch_size, 1, self.subgraph_size)).to(self.device)
+        added_adj_zero_col = torch.zeros(
+            (batch_size, self.subgraph_size + 1, 1)).to(self.device)
+        added_adj_zero_col[:, -1, :] = 1.
+        added_feat_zero_row = torch.zeros((batch_size, 1,
+                                           self.n_in)).to(self.device)
+
+        for i in idx:
+            cur_adj = adj[:, subgraphs[i], :][:, :, subgraphs[i]]
+            cur_feat = x[:, subgraphs[i], :]
+            batch_adj.append(cur_adj)
+            batch_feature.append(cur_feat)
+
+        batch_adj = torch.cat(batch_adj)
+        batch_adj = torch.cat((batch_adj, added_adj_zero_row), dim=1)
+        batch_adj = torch.cat((batch_adj, added_adj_zero_col), dim=2)
+        batch_feature = torch.cat(batch_feature)
+        batch_feature = torch.cat(
+            (batch_feature[:, :-1, :],
+             added_feat_zero_row,
+             batch_feature[:, -1:, :]), dim=1)
+
+        h_1 = self.gcn(batch_feature, batch_adj, sparse)
+
+        if self.readout == 'max':
+            h_mv = h_1[:, -1, :]
+            c = torch.max(h_1[:, : -1, :], 1).values
+        elif self.readout == 'min':
+            h_mv = h_1[:, -1, :]
+            c = torch.min(h_1[:, : -1, :], 1).values
+        elif self.readout == 'avg':
+            h_mv = h_1[:, -1, :]
+            c = torch.mean(h_1[:, : -1, :], 1)
+        elif self.readout == 'weighted_sum':
+            seq, query = h_1[:, : -1, :], h_1[:, -2: -1, :],
+            query = query.permute(0, 2, 1)
+            sim = torch.matmul(seq, query)
+            sim = F.softmax(sim, dim=1)
+            sim = sim.repeat(1, 1, 64)
+            out = torch.mul(seq, sim)
+            c = torch.sum(out, 1)
+            h_mv = h_1[:, -1, :]
+
+        ret = self.disc(c, h_mv)
+
+        return ret
+
+
+class Discriminator(nn.Module):
+    def __init__(self, n_h, negsamp_ratio):
+        super(Discriminator, self).__init__()
+        self.f_k = nn.Bilinear(n_h, n_h, 1)
+
+        for m in self.modules():
+            self.weights_init(m)
+
+        self.negsamp_ratio = negsamp_ratio
+
+    def weights_init(self, m):
+        if isinstance(m, nn.Bilinear):
+            torch.nn.init.xavier_uniform_(m.weight.data)
+            if m.bias is not None:
+                m.bias.data.fill_(0.0)
+
+    def forward(self, c, h_pl):
+        scs = []
+        # positive
+        scs.append(self.f_k(h_pl, c))
+
+        # negative
+        c_mi = c
+        for _ in range(self.negsamp_ratio):
+            c_mi = torch.cat((c_mi[-2:-1, :], c_mi[:-1, :]), 0)
+            scs.append(self.f_k(h_pl, c_mi))
+
+        logits = torch.cat(tuple(scs))
+
+        return logits
diff --git a/pygod/models/conad.py b/pygod/models/conad.py
new file mode 100644
index 0000000..4a53208
--- /dev/null
+++ b/pygod/models/conad.py
@@ -0,0 +1,422 @@
+# -*- coding: utf-8 -*-
+"""Contrastive Attributed Network Anomaly Detection
+with Data Augmentation (CONAD)"""
+# Author: Zhiming Xu <zhimng.xu@gmail.com>
+# License: BSD 2 clause
+
+import torch
+import numpy as np
+import torch.nn as nn
+from copy import deepcopy
+import torch.nn.functional as F
+from torch_geometric.loader import NeighborLoader
+from torch_geometric.utils import to_dense_adj, dense_to_sparse
+from sklearn.utils.validation import check_is_fitted
+
+from . import BaseDetector
+from .basic_nn import GCN
+from ..utils.metric import eval_roc_auc
+from ..utils.utility import validate_device
+
+
+class CONAD(BaseDetector):
+    """
+    CONAD (Contrastive Attributed Network Anomaly Detection)
+    CONAD is an anomaly detector consisting of a shared graph
+    convolutional encoder, a structure reconstruction decoder, and an
+    attribute reconstruction decoder. The model is trained with both
+    contrastive loss and structure/attribute reconstruction loss.
+    The reconstruction mean square error of the decoders are defined
+    as structure anomaly score and attribute anomaly score, respectively.
+
+    See :cite:`xu2022contrastive` for details.
+
+    Parameters
+    ----------
+    hid_dim :  int, optional
+        Hidden dimension of model. Default: ``0``.
+    num_layers : int, optional
+        Total number of layers in model. A half (ceil) of the layers
+        are for the encoder, the other half (floor) of the layers are
+        for decoders. Default: ``4``.
+    dropout : float, optional
+        Dropout rate. Default: ``0.``.
+    weight_decay : float, optional
+        Weight decay (L2 penalty). Default: ``0.``.
+    act : callable activation function or None, optional
+        Activation function if not None.
+        Default: ``torch.nn.functional.relu``.
+    alpha : float, optional
+        Loss balance weight for attribute and structure.
+        Default: ``0.5``.
+    eta : float, optional
+        Loss balance weight for contrastive and reconstruction.
+        Default: ``0.5``.
+    contamination : float, optional
+        Valid in (0., 0.5). The proportion of outliers in the data set.
+        Used when fitting to define the threshold on the decision
+        function. Default: ``0.1``.
+    lr : float, optional
+        Learning rate. Default: ``0.004``.
+    epoch : int, optional
+        Maximum number of training epoch. Default: ``5``.
+    gpu : int
+        GPU Index, -1 for using CPU. Default: ``0``.
+    batch_size : int, optional
+        Minibatch size, 0 for full batch training. Default: ``0``.
+    num_neigh : int, optional
+        Number of neighbors in sampling, -1 for all neighbors.
+        Default: ``-1``.
+    r : float, optional
+        The rate of augmented anomalies. Default: ``.2``.
+    m : int, optional
+        For densely connected nodes, the number of
+        edges to add. Default: ``50``.
+    k : int, optional
+        same as ``k`` in ``utils.outlier_generator.gen_attribute_outliers``.
+        Default: ``50``.
+    f : int, optional
+        For disproportionate nodes, the scale factor applied
+        on their attribute value. Default: ``10``.
+    verbose : bool
+        Verbosity mode. Turn on to print out log information.
+        Default: ``False``.
+
+    Examples
+    --------
+    >>> from pygod.models import CONAD
+    >>> model = CONAD()
+    >>> model.fit(data) # PyG graph data object
+    >>> prediction = model.predict(data)
+    """
+
+    def __init__(self,
+                 hid_dim=64,
+                 num_layers=4,
+                 dropout=0.3,
+                 weight_decay=0.,
+                 act=F.relu,
+                 alpha=0.8,
+                 eta=.5,
+                 contamination=0.1,
+                 lr=5e-3,
+                 epoch=5,
+                 gpu=0,
+                 batch_size=0,
+                 num_neigh=-1,
+                 margin=.5,
+                 r=.2,
+                 m=50,
+                 k=50,
+                 f=10,
+                 verbose=False):
+        super(CONAD, self).__init__(contamination=contamination)
+
+        # model param
+        self.hid_dim = hid_dim
+        self.num_layers = num_layers
+        self.dropout = dropout
+        self.weight_decay = weight_decay
+        self.act = act
+        self.alpha = alpha
+        self.eta = eta
+
+        # training param
+        self.lr = lr
+        self.epoch = epoch
+        self.device = validate_device(gpu)
+        self.batch_size = batch_size
+        self.num_neigh = num_neigh
+
+        self.margin_loss_func = torch.nn.MarginRankingLoss(margin=margin)
+        # other param
+        self.verbose = verbose
+        self.r = r
+        self.m = m
+        self.k = k
+        self.f = f
+        self.model = None
+
+    def fit(self, G, y_true=None):
+        """
+        Fit detector with input data.
+
+        Parameters
+        ----------
+        G : PyTorch Geometric Data instance (torch_geometric.data.Data)
+            The input data.
+        y_true : numpy.ndarray, optional
+            The optional outlier ground truth labels used to monitor
+            the training progress. They are not used to optimize the
+            unsupervised model. Default: ``None``.
+
+        Returns
+        -------
+        self : object
+            Fitted estimator.
+        """
+        G.node_idx = torch.arange(G.x.shape[0])
+        G.s = to_dense_adj(G.edge_index)[0]
+        if self.batch_size == 0:
+            self.batch_size = G.x.shape[0]
+        loader = NeighborLoader(G,
+                                [self.num_neigh] * self.num_layers,
+                                batch_size=self.batch_size)
+
+        self.model = CONAD_Base(in_dim=G.x.shape[1],
+                                hid_dim=self.hid_dim,
+                                num_layers=self.num_layers,
+                                dropout=self.dropout,
+                                act=self.act).to(self.device)
+
+        optimizer = torch.optim.Adam(self.model.parameters(),
+                                     lr=self.lr,
+                                     weight_decay=self.weight_decay)
+
+        self.model.train()
+        decision_scores = np.zeros(G.x.shape[0])
+        for epoch in range(self.epoch):
+            epoch_loss = 0
+            for sampled_data in loader:
+                batch_size = sampled_data.batch_size
+                node_idx = sampled_data.node_idx
+                x, s, edge_index = self.process_graph(sampled_data)
+
+                # generate augmented graph
+                x_aug, edge_index_aug, label_aug = \
+                    self._data_augmentation(x, s)
+                h_aug = self.model.embed(x_aug, edge_index_aug)
+                h = self.model.embed(x, edge_index)
+
+                margin_loss = self.margin_loss_func(h, h, h_aug) * label_aug
+                margin_loss = torch.mean(margin_loss)
+
+                x_, s_ = self.model.reconstruct(h, edge_index)
+                score = self.loss_func(x, x_, s, s_)
+                loss = self.eta * torch.mean(score) + \
+                    (1 - self.eta) * margin_loss
+
+                decision_scores[node_idx[:batch_size]] = score.detach() \
+                    .cpu().numpy()
+                epoch_loss += loss.item() * batch_size
+
+                optimizer.zero_grad()
+                loss.backward()
+                optimizer.step()
+
+            if self.verbose:
+                print("Epoch {:04d}: Loss {:.4f}"
+                      .format(epoch, epoch_loss / G.x.shape[0]), end='')
+                if y_true is not None:
+                    auc = eval_roc_auc(y_true, decision_scores)
+                    print(" | AUC {:.4f}".format(auc), end='')
+                print()
+
+        self.decision_scores_ = decision_scores
+        self._process_decision_scores()
+        return self
+
+    def decision_function(self, G):
+        """
+        Predict raw anomaly score using the fitted detector. Outliers
+        are assigned with larger anomaly scores.
+
+        Parameters
+        ----------
+        G : PyTorch Geometric Data instance (torch_geometric.data.Data)
+            The input data.
+
+        Returns
+        -------
+        outlier_scores : numpy.ndarray
+            The anomaly score of shape :math:`N`.
+        """
+        check_is_fitted(self, ['model'])
+        G.node_idx = torch.arange(G.x.shape[0])
+        G.s = to_dense_adj(G.edge_index)[0]
+
+        loader = NeighborLoader(G,
+                                [self.num_neigh] * self.num_layers,
+                                batch_size=self.batch_size)
+
+        self.model.eval()
+        outlier_scores = np.zeros(G.x.shape[0])
+        for sampled_data in loader:
+            batch_size = sampled_data.batch_size
+            node_idx = sampled_data.node_idx
+
+            x, s, edge_index = self.process_graph(sampled_data)
+
+            x_, s_ = self.model(x, edge_index)
+            score = self.loss_func(x[:batch_size],
+                                   x_[:batch_size],
+                                   s[:batch_size],
+                                   s_[:batch_size])
+
+            outlier_scores[node_idx[:batch_size]] = score.detach() \
+                .cpu().numpy()
+        return outlier_scores
+
+    def _data_augmentation(self, x, adj):
+        """
+        Data augmentation on the input graph. Four types of
+        pseudo anomalies will be injected:
+            Attribute, deviated
+            Attribute, disproportionate
+            Structure, high-degree
+            Structure, outlying
+        
+        Parameters
+        -----------
+        x : note attribute matrix
+        adj : dense adjacency matrix
+
+        Returns
+        -------
+        feat_aug, adj_aug, label_aug : augmented
+            attribute matrix, adjacency matrix, and
+            pseudo anomaly label to train contrastive
+            graph representations
+        """
+        rate = self.r
+        num_added_edge = self.m
+        surround = self.k
+        scale_factor = self.f
+
+        adj_aug, feat_aug = deepcopy(adj), deepcopy(x)
+        num_nodes = adj_aug.shape[0]
+        label_aug = torch.zeros(num_nodes, dtype=torch.int32)
+
+        for i in range(num_nodes):
+            prob = np.random.uniform()
+            if prob > rate: continue
+            label_aug[i] = 1
+            one_fourth = np.random.randint(0, 4)
+            if one_fourth == 0:
+                # add clique
+                new_neighbors = np.random.choice(np.arange(num_nodes),
+                                                 num_added_edge, replace=False)
+                adj_aug[i, new_neighbors] = 1
+                adj_aug[new_neighbors, i] = 1
+            elif one_fourth == 1:
+                # drop all connection
+                neighbors = torch.nonzero(adj[i]).view(-1)
+                if not neighbors.any():
+                    continue
+                adj_aug[i, neighbors] = 0
+                adj_aug[neighbors, i] = 0
+            elif one_fourth == 2:
+                # attrs
+                candidates = np.random.choice(np.arange(num_nodes), surround,
+                                              replace=False)
+                max_dev, max_idx = 0, i
+                for c in candidates:
+                    dev = torch.square(feat_aug[i] - feat_aug[c]).sum()
+                    if dev > max_dev:
+                        max_dev = dev
+                        max_idx = c
+                feat_aug[i] = feat_aug[max_idx]
+            else:
+                # scale attr
+                prob = np.random.uniform(0, 1)
+                if prob > 0.5:
+                    feat_aug[i] *= scale_factor
+                else:
+                    feat_aug[i] /= scale_factor
+
+        edge_index_aug = dense_to_sparse(adj_aug)[0].to(self.device)
+        feat_aug = feat_aug.to(self.device)
+        label_aug = label_aug.to(device=self.device)
+        return feat_aug, edge_index_aug, label_aug
+
+    def process_graph(self, G):
+        """
+        Process the raw PyG data object into a tuple of sub data
+        objects needed for the model.
+
+        Parameters
+        ----------
+        G : PyTorch Geometric Data instance (torch_geometric.data.Data)
+            The input data.
+
+        Returns
+        -------
+        x : torch.Tensor
+            Attribute (feature) of nodes.
+        s : torch.Tensor
+            Adjacency matrix of the graph.
+        edge_index : torch.Tensor
+            Edge list of the graph.
+        """
+        s = G.s.to(self.device)
+        edge_index = G.edge_index.to(self.device)
+        x = G.x.to(self.device)
+
+        return x, s, edge_index
+
+    def loss_func(self, x, x_, s, s_):
+        # attribute reconstruction loss
+        diff_attribute = torch.pow(x - x_, 2)
+        attribute_errors = torch.sqrt(torch.sum(diff_attribute, 1))
+
+        # structure reconstruction loss
+        diff_structure = torch.pow(s - s_, 2)
+        structure_errors = torch.sqrt(torch.sum(diff_structure, 1))
+
+        score = self.alpha * attribute_errors \
+                + (1 - self.alpha) * structure_errors
+        return score
+
+
+class CONAD_Base(nn.Module):
+    def __init__(self,
+                 in_dim,
+                 hid_dim,
+                 num_layers,
+                 dropout,
+                 act):
+        super(CONAD_Base, self).__init__()
+
+        decoder_layers = int(num_layers / 2)
+        encoder_layers = num_layers - decoder_layers
+
+        self.shared_encoder = GCN(in_channels=in_dim,
+                                  hidden_channels=hid_dim,
+                                  num_layers=encoder_layers,
+                                  out_channels=hid_dim,
+                                  dropout=dropout,
+                                  act=act)
+
+        self.attr_decoder = GCN(in_channels=hid_dim,
+                                hidden_channels=hid_dim,
+                                num_layers=decoder_layers,
+                                out_channels=in_dim,
+                                dropout=dropout,
+                                act=act)
+
+        self.struct_decoder = GCN(in_channels=hid_dim,
+                                  hidden_channels=hid_dim,
+                                  num_layers=decoder_layers - 1,
+                                  out_channels=in_dim,
+                                  dropout=dropout,
+                                  act=act)
+
+    def embed(self, x, edge_index):
+        h = self.shared_encoder(x, edge_index)
+        return h
+
+    def reconstruct(self, h, edge_index):
+        # decode attribute matrix
+        x_ = self.attr_decoder(h, edge_index)
+        # decode structure matrix
+        h_ = self.struct_decoder(h, edge_index)
+
+        s_ = h_ @ h_.T
+        return x_, s_
+
+    def forward(self, x, edge_index):
+        # encode
+        h = self.embed(x, edge_index)
+        # reconstruct
+        x_, s_ = self.reconstruct(h, edge_index)
+        return x_, s_
diff --git a/pygod/models/dominant.py b/pygod/models/dominant.py
index 55038b1..12bc313 100644
--- a/pygod/models/dominant.py
+++ b/pygod/models/dominant.py
@@ -3,14 +3,17 @@
 # Author: Kay Liu <zliu234@uic.edu>
 # License: BSD 2 clause
 
+import numpy as np
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
 from torch_geometric.utils import to_dense_adj
+from torch_geometric.loader import NeighborLoader
 from sklearn.utils.validation import check_is_fitted
 
 from . import BaseDetector
 from .basic_nn import GCN
+from ..utils.utility import validate_device
 from ..utils.metric import eval_roc_auc
 
 
@@ -53,6 +56,11 @@ class DOMINANT(BaseDetector):
         Maximum number of training epoch. Default: ``5``.
     gpu : int
         GPU Index, -1 for using CPU. Default: ``0``.
+    batch_size : int, optional
+        Minibatch size, 0 for full batch training. Default: ``0``.
+    num_neigh : int, optional
+        Number of neighbors in sampling, -1 for all neighbors.
+        Default: ``-1``.
     verbose : bool
         Verbosity mode. Turn on to print out log information.
         Default: ``False``.
@@ -64,6 +72,7 @@ class DOMINANT(BaseDetector):
     >>> model.fit(data) # PyG graph data object
     >>> prediction = model.predict(data)
     """
+
     def __init__(self,
                  hid_dim=64,
                  num_layers=4,
@@ -75,6 +84,8 @@ def __init__(self,
                  lr=5e-3,
                  epoch=5,
                  gpu=0,
+                 batch_size=0,
+                 num_neigh=-1,
                  verbose=False):
         super(DOMINANT, self).__init__(contamination=contamination)
 
@@ -89,10 +100,9 @@ def __init__(self,
         # training param
         self.lr = lr
         self.epoch = epoch
-        if gpu >= 0 and torch.cuda.is_available():
-            self.device = 'cuda:{}'.format(gpu)
-        else:
-            self.device = 'cpu'
+        self.device = validate_device(gpu)
+        self.batch_size = batch_size
+        self.num_neigh = num_neigh
 
         # other param
         self.verbose = verbose
@@ -100,28 +110,31 @@ def __init__(self,
 
     def fit(self, G, y_true=None):
         """
-        Description
-        -----------
         Fit detector with input data.
 
         Parameters
         ----------
-        G : PyTorch Geometric Data instance (torch_geometric.data.Data)
+        G : torch_geometric.data.Data
             The input data.
-        y_true : numpy.array, optional (default=None)
-            The optional outlier ground truth labels used to monitor the
-            training progress. They are not used to optimize the
-            unsupervised model.
+        y_true : numpy.ndarray, optional
+            The optional outlier ground truth labels used to monitor
+            the training progress. They are not used to optimize the
+            unsupervised model. Default: ``None``.
 
         Returns
         -------
         self : object
             Fitted estimator.
         """
-
-        x, adj, edge_index = self.process_graph(G)
-
-        self.model = DOMINANT_Base(in_dim=x.shape[1],
+        G.node_idx = torch.arange(G.x.shape[0])
+        G.s = to_dense_adj(G.edge_index)[0]
+        if self.batch_size == 0:
+            self.batch_size = G.x.shape[0]
+        loader = NeighborLoader(G,
+                                [self.num_neigh] * self.num_layers,
+                                batch_size=self.batch_size)
+
+        self.model = DOMINANT_Base(in_dim=G.x.shape[1],
                                    hid_dim=self.hid_dim,
                                    num_layers=self.num_layers,
                                    dropout=self.dropout,
@@ -130,33 +143,44 @@ def fit(self, G, y_true=None):
         optimizer = torch.optim.Adam(self.model.parameters(),
                                      lr=self.lr,
                                      weight_decay=self.weight_decay)
-        score = None
-        for epoch in range(self.epoch):
-            self.model.train()
-            x_, adj_ = self.model(x, edge_index)
-            score = self.loss_func(x, x_, adj, adj_)
-            loss = torch.mean(score)
 
-            optimizer.zero_grad()
-            loss.backward()
-            optimizer.step()
+        self.model.train()
+        decision_scores = np.zeros(G.x.shape[0])
+        for epoch in range(self.epoch):
+            epoch_loss = 0
+            for sampled_data in loader:
+                batch_size = sampled_data.batch_size
+                node_idx = sampled_data.node_idx
+                x, s, edge_index = self.process_graph(sampled_data)
+
+                x_, s_ = self.model(x, edge_index)
+                score = self.loss_func(x[:batch_size],
+                                       x_[:batch_size],
+                                       s[:batch_size, node_idx],
+                                       s_[:batch_size])
+                decision_scores[node_idx[:batch_size]] = score.detach() \
+                    .cpu().numpy()
+                loss = torch.mean(score)
+                epoch_loss += loss.item() * batch_size
+
+                optimizer.zero_grad()
+                loss.backward()
+                optimizer.step()
 
             if self.verbose:
                 print("Epoch {:04d}: Loss {:.4f}"
-                      .format(epoch, loss.item()), end='')
+                      .format(epoch, epoch_loss / G.x.shape[0]), end='')
                 if y_true is not None:
-                    auc = eval_roc_auc(y_true, score.detach().cpu().numpy())
+                    auc = eval_roc_auc(y_true, decision_scores)
                     print(" | AUC {:.4f}".format(auc), end='')
                 print()
 
-        self.decision_scores_ = score.detach().cpu().numpy()
+        self.decision_scores_ = decision_scores
         self._process_decision_scores()
         return self
 
     def decision_function(self, G):
         """
-        Description
-        -----------
         Predict raw anomaly score using the fitted detector. Outliers
         are assigned with larger anomaly scores.
 
@@ -171,22 +195,33 @@ def decision_function(self, G):
             The anomaly score of shape :math:`N`.
         """
         check_is_fitted(self, ['model'])
+        G.node_idx = torch.arange(G.x.shape[0])
+        G.s = to_dense_adj(G.edge_index)[0]
 
-        # get needed data object from the input data
-        x, adj, edge_index = self.process_graph(G)
+        loader = NeighborLoader(G,
+                                [self.num_neigh] * self.num_layers,
+                                batch_size=self.batch_size)
 
-        # enable the evaluation mode
         self.model.eval()
+        outlier_scores = np.zeros(G.x.shape[0])
+        for sampled_data in loader:
+            batch_size = sampled_data.batch_size
+            node_idx = sampled_data.node_idx
 
-        # construct the vector for holding the reconstruction error
-        x_, adj_ = self.model(x, edge_index)
-        outlier_scores = self.loss_func(x, x_, adj, adj_)
-        return outlier_scores.detach().cpu().numpy()
+            x, s, edge_index = self.process_graph(sampled_data)
+
+            x_, s_ = self.model(x, edge_index)
+            score = self.loss_func(x[:batch_size],
+                                   x_[:batch_size],
+                                   s[:batch_size, node_idx],
+                                   s_[:batch_size])
+
+            outlier_scores[node_idx[:batch_size]] = score.detach() \
+                .cpu().numpy()
+        return outlier_scores
 
     def process_graph(self, G):
         """
-        Description
-        -----------
         Process the raw PyG data object into a tuple of sub data
         objects needed for the model.
 
@@ -199,29 +234,24 @@ def process_graph(self, G):
         -------
         x : torch.Tensor
             Attribute (feature) of nodes.
-        adj : torch.Tensor
+        s : torch.Tensor
             Adjacency matrix of the graph.
         edge_index : torch.Tensor
             Edge list of the graph.
         """
-        edge_index = G.edge_index
-
-        adj = to_dense_adj(edge_index)[0].to(self.device)
-
-        edge_index = edge_index.to(self.device)
-        adj = adj.to(self.device)
+        s = G.s.to(self.device)
+        edge_index = G.edge_index.to(self.device)
         x = G.x.to(self.device)
 
-        # return data objects needed for the network
-        return x, adj, edge_index
+        return x, s, edge_index
 
-    def loss_func(self, x, x_, adj, adj_):
+    def loss_func(self, x, x_, s, s_):
         # attribute reconstruction loss
         diff_attribute = torch.pow(x - x_, 2)
         attribute_errors = torch.sqrt(torch.sum(diff_attribute, 1))
 
         # structure reconstruction loss
-        diff_structure = torch.pow(adj - adj_, 2)
+        diff_structure = torch.pow(s - s_, 2)
         structure_errors = torch.sqrt(torch.sum(diff_structure, 1))
 
         score = self.alpha * attribute_errors \
@@ -236,7 +266,6 @@ def __init__(self,
                  num_layers,
                  dropout,
                  act):
-
         super(DOMINANT_Base, self).__init__()
 
         # split the number of layers for the encoder and decoders
@@ -271,7 +300,6 @@ def forward(self, x, edge_index):
         x_ = self.attr_decoder(h, edge_index)
         # decode adjacency matrix
         h_ = self.struct_decoder(h, edge_index)
-        adj_ = h_ @ h_.T
-
+        s_ = h_ @ h_.T
         # return reconstructed matrices
-        return x_, adj_
+        return x_, s_
diff --git a/pygod/models/done.py b/pygod/models/done.py
index 6c5fd6e..19ef9d0 100644
--- a/pygod/models/done.py
+++ b/pygod/models/done.py
@@ -4,21 +4,24 @@
 # License: BSD 2 clause
 
 import torch
+import numpy as np
 import torch.nn as nn
 import torch.nn.functional as F
-from torch_geometric.utils import to_dense_adj
 from torch_geometric.nn import MessagePassing
+from torch_geometric.utils import to_dense_adj
+from torch_geometric.loader import NeighborLoader
 from sklearn.utils.validation import check_is_fitted
 
 from . import BaseDetector
 from .basic_nn import MLP
+from ..utils.utility import validate_device
 from ..utils.metric import eval_roc_auc
 
 
 class DONE(BaseDetector):
     """
     DONE (Deep Outlier Aware Attributed Network Embedding)
-    DONE is consist of an attribute autoencoder and a structure
+    DONE consist of an attribute autoencoder and a structure
     autoencoder. It estimates five loss to optimize the model,
     including an attribute proximity loss, an attribute homophily loss,
     a structure proximity loss, a structure homophily loss, and a
@@ -68,6 +71,11 @@ class DONE(BaseDetector):
         Maximum number of training epoch. Default: ``5``.
     gpu : int
         GPU Index, -1 for using CPU. Default: ``0``.
+    batch_size : int, optional
+        Minibatch size, 0 for full batch training. Default: ``0``.
+    num_neigh : int, optional
+        Number of neighbors in sampling, -1 for all neighbors.
+        Default: ``-1``.
     verbose : bool
         Verbosity mode. Turn on to print out log information.
         Default: ``False``.
@@ -95,6 +103,8 @@ def __init__(self,
                  lr=5e-3,
                  epoch=5,
                  gpu=0,
+                 batch_size=0,
+                 num_neigh=-1,
                  verbose=False):
         super(DONE, self).__init__(contamination=contamination)
 
@@ -113,10 +123,9 @@ def __init__(self,
         # training param
         self.lr = lr
         self.epoch = epoch
-        if gpu >= 0 and torch.cuda.is_available():
-            self.device = 'cuda:{}'.format(gpu)
-        else:
-            self.device = 'cpu'
+        self.device = validate_device(gpu)
+        self.batch_size = batch_size
+        self.num_neigh = num_neigh
 
         # other param
         self.verbose = verbose
@@ -124,29 +133,32 @@ def __init__(self,
 
     def fit(self, G, y_true=None):
         """
-        Description
-        -----------
         Fit detector with input data.
 
         Parameters
         ----------
-        G : PyTorch Geometric Data instance (torch_geometric.data.Data)
+        G : torch_geometric.data.Data
             The input data.
-        y_true : numpy.array, optional (default=None)
-            The optional outlier ground truth labels used to monitor the
-            training progress. They are not used to optimize the
-            unsupervised model.
+        y_true : numpy.ndarray, optional
+            The optional outlier ground truth labels used to monitor
+            the training progress. They are not used to optimize the
+            unsupervised model. Default: ``None``.
 
         Returns
         -------
         self : object
             Fitted estimator.
         """
-
-        x, s, edge_index = self.process_graph(G)
-
-        self.model = DONE_Base(x_dim=x.shape[1],
-                               s_dim=s.shape[1],
+        G.node_idx = torch.arange(G.x.shape[0])
+        G.s = to_dense_adj(G.edge_index)[0]
+        if self.batch_size == 0:
+            self.batch_size = G.x.shape[0]
+        loader = NeighborLoader(G,
+                                [self.num_neigh],
+                                batch_size=self.batch_size)
+
+        self.model = DONE_Base(x_dim=G.x.shape[1],
+                               s_dim=G.s.shape[1],
                                hid_dim=self.hid_dim,
                                num_layers=self.num_layers,
                                dropout=self.dropout,
@@ -155,32 +167,48 @@ def fit(self, G, y_true=None):
         optimizer = torch.optim.Adam(self.model.parameters(),
                                      lr=self.lr,
                                      weight_decay=self.weight_decay)
-        score = None
-        for epoch in range(self.epoch):
-            self.model.train()
-            x_, s_, h_a, h_s, dna, dns = self.model(x, s, edge_index)
-            score, loss = self.loss_func(x, x_, s, s_, h_a, h_s, dna, dns)
 
-            optimizer.zero_grad()
-            loss.backward()
-            optimizer.step()
+        self.model.train()
+        decision_scores = np.zeros(G.x.shape[0])
+        for epoch in range(self.epoch):
+            epoch_loss = 0
+            for sampled_data in loader:
+                batch_size = sampled_data.batch_size
+                node_idx = sampled_data.node_idx
+                x, s, edge_index = self.process_graph(G)
+
+                x_, s_, h_a, h_s, dna, dns = self.model(x, s, edge_index)
+                score, loss = self.loss_func(x[:batch_size],
+                                             x_[:batch_size],
+                                             s[:batch_size],
+                                             s_[:batch_size],
+                                             h_a[:batch_size],
+                                             h_s[:batch_size],
+                                             dna[:batch_size],
+                                             dns[:batch_size])
+                epoch_loss += loss.item() * batch_size
+                decision_scores[node_idx[:batch_size]] = score.detach() \
+                                                              .cpu().numpy()
+
+                optimizer.zero_grad()
+                loss.backward()
+                optimizer.step()
 
             if self.verbose:
                 print("Epoch {:04d}: Loss {:.4f}"
-                      .format(epoch, loss.item()), end='')
+                      .format(epoch, epoch_loss / G.x.shape[0]), end='')
                 if y_true is not None:
-                    auc = eval_roc_auc(y_true, score.detach().cpu().numpy())
+                    auc = eval_roc_auc(y_true, decision_scores)
                     print(" | AUC {:.4f}".format(auc), end='')
                 print()
 
-        self.decision_scores_ = score.detach().cpu().numpy()
+        self.decision_scores_ = decision_scores
         self._process_decision_scores()
         return self
 
     def decision_function(self, G):
         """
-        Description
-        -----------
+
         Predict raw anomaly score using the fitted detector. Outliers
         are assigned with larger anomaly scores.
 
@@ -195,22 +223,38 @@ def decision_function(self, G):
             The anomaly score of shape :math:`N`.
         """
         check_is_fitted(self, ['model'])
+        G.node_idx = torch.arange(G.x.shape[0])
+        G.s = to_dense_adj(G.edge_index)[0]
+        if self.batch_size == 0:
+            self.batch_size = G.x.shape[0]
+        loader = NeighborLoader(G,
+                                [self.num_neigh],
+                                batch_size=self.batch_size)
 
-        # get needed data object from the input data
-        x, s, edge_index = self.process_graph(G)
-
-        # enable the evaluation mode
         self.model.eval()
+        outlier_scores = np.zeros(G.x.shape[0])
+        for sampled_data in loader:
+            batch_size = sampled_data.batch_size
+            node_idx = sampled_data.node_idx
 
-        # construct the vector for holding the reconstruction error
-        x_, s_, h_a, h_s, dna, dns = self.model(x, s, edge_index)
-        outlier_scores, _ = self.loss_func(x, x_, s, s_, h_a, h_s, dna, dns)
-        return outlier_scores.detach().cpu().numpy()
+            x, s, edge_index = self.process_graph(G)
+
+            x_, s_, h_a, h_s, dna, dns = self.model(x, s, edge_index)
+            score, _ = self.loss_func(x[:batch_size],
+                                      x_[:batch_size],
+                                      s[:batch_size],
+                                      s_[:batch_size],
+                                      h_a[:batch_size],
+                                      h_s[:batch_size],
+                                      dna[:batch_size],
+                                      dns[:batch_size])
+
+            outlier_scores[node_idx[:batch_size]] = score.detach() \
+                                                         .cpu().numpy()
+        return outlier_scores
 
     def process_graph(self, G):
         """
-        Description
-        -----------
         Process the raw PyG data object into a tuple of sub data
         objects needed for the model.
 
@@ -228,10 +272,8 @@ def process_graph(self, G):
         edge_index : torch.Tensor
             Edge list of the graph.
         """
-        edge_index = G.edge_index
-
-        s = to_dense_adj(edge_index)[0].to(self.device)
-        edge_index = edge_index.to(self.device)
+        s = G.s.to(self.device)
+        edge_index = G.edge_index.to(self.device)
         x = G.x.to(self.device)
 
         return x, s, edge_index
diff --git a/pygod/models/gaan.py b/pygod/models/gaan.py
index 2375911..72e2aaf 100644
--- a/pygod/models/gaan.py
+++ b/pygod/models/gaan.py
@@ -3,12 +3,16 @@
 # Author: Ruitong Zhang <rtzhang@buaa.edu.cn>
 # License: BSD 2 clause
 
+import numpy as np
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
+from torch_geometric.utils import to_dense_adj
+from torch_geometric.loader import NeighborLoader
 from sklearn.utils.validation import check_is_fitted
 
 from . import BaseDetector
+from ..utils.utility import validate_device
 from ..utils.metric import eval_roc_auc
 
 
@@ -36,9 +40,7 @@ class GAAN(BaseDetector):
     hid_dim3 :  int, optional
         Hidden dimension of MLP later 3. Defaults: ``128``.
     num_layers : int, optional
-        Total number of layers in model. A half (ceil) of the layers
-        are for the encoder, the other half (floor) of the layers are
-        for decoders. Defaults: ``3``.
+        Total number of layers in model. Defaults: ``3``.
     dropout : float, optional
         Dropout rate. Defaults: ``0.``.
     weight_decay : float, optional
@@ -59,6 +61,11 @@ class GAAN(BaseDetector):
         Maximum number of training epoch. Defaults: ``10``.
     gpu : int
         GPU Index, -1 for using CPU. Defaults: ``-1``.
+    batch_size : int, optional
+        Minibatch size, 0 for full batch training. Default: ``0``.
+    num_neigh : int, optional
+        Number of neighbors in sampling, -1 for all neighbors.
+        Default: ``-1``.
     verbose : bool
         Verbosity mode. Turn on to print out log information.
         Defaults: ``False``.
@@ -77,7 +84,7 @@ def __init__(self,
                  hid_dim1=32,
                  hid_dim2=64,
                  hid_dim3=128,
-                 num_layers=3,
+                 num_layers=2,
                  dropout=0.3,
                  weight_decay=0.,
                  act=F.relu,
@@ -85,7 +92,9 @@ def __init__(self,
                  contamination=0.1,
                  lr=5e-3,
                  epoch=10,
-                 gpu=-1,
+                 gpu=0,
+                 batch_size=0,
+                 num_neigh=-1,
                  verbose=False):
         super(GAAN, self).__init__(contamination=contamination)
 
@@ -104,10 +113,9 @@ def __init__(self,
         # training param
         self.lr = lr
         self.epoch = epoch
-        if gpu >= 0 and torch.cuda.is_available():
-            self.device = 'cuda:{}'.format(gpu)
-        else:
-            self.device = 'cpu'
+        self.device = validate_device(gpu)
+        self.batch_size = batch_size
+        self.num_neigh = num_neigh
 
         # other param
         self.verbose = verbose
@@ -119,35 +127,40 @@ def __init__(self,
 
     def fit(self, G, y_true=None):
         """
-        Description
-        -----------
         Fit detector with input data.
 
         Parameters
         ----------
-        G : PyTorch Geometric Data instance (torch_geometric.data.Data)
+        G : torch_geometric.data.Data
             The input data.
-        y_true : numpy.array, optional (default=None)
-            The optional outlier ground truth labels used to monitor the
-            training progress. They are not used to optimize the
-            unsupervised model.
+        y_true : numpy.ndarray, optional
+            The optional outlier ground truth labels used to monitor
+            the training progress. They are not used to optimize the
+            unsupervised model. Default: ``None``.
 
         Returns
         -------
         self : object
             Fitted estimator.
         """
-        X, edge_index = self.process_graph(G)
+        #X, edge_index = self.process_graph(G)
+        G.node_idx = torch.arange(G.x.shape[0])
+        G.s = to_dense_adj(G.edge_index)[0]
+        if self.batch_size == 0:
+            self.batch_size = G.x.shape[0]
+        loader = NeighborLoader(G,
+                                [self.num_neigh] * self.num_layers,
+                                batch_size=self.batch_size)
 
         # initialize the model
         self.generator = Generator(in_dim=self.noise_dim,
                                    hid_dim1=self.hid_dim1,
                                    hid_dim2=self.hid_dim2,
                                    hid_dim3=self.hid_dim3,
-                                   out_dim=X.shape[1],
+                                   out_dim=G.x.shape[1],
                                    act=self.act).to(self.device)
 
-        self.encoder = Encoder(in_dim=X.shape[1],
+        self.encoder = Encoder(in_dim=G.x.shape[1],
                                hid_dim1=self.hid_dim1,
                                hid_dim2=self.hid_dim2,
                                hid_dim3=self.hid_dim3,
@@ -166,59 +179,82 @@ def fit(self, G, y_true=None):
         optimizer_D = torch.optim.Adam(params=self.discriminator.parameters(),
                                        weight_decay=self.weight_decay)
 
+        # enable the training mode
+        self.generator.train()
+        self.encoder.train()
+        self.discriminator.train()
+
         # initialize the  criterion
         criterion = torch.nn.BCELoss()
-        X_, Y_true_pre, Y_fake_pre = None, None, None
 
+        scores = np.zeros(G.x.shape[0])
         for epoch in range(self.epoch):
-
-            # Generate noise for constructing fake attribute
-            gaussian_noise = torch.randn(X.shape[0], self.noise_dim).to(
-                self.device)
-
-            self.generator.train()
-            self.encoder.train()
-            self.discriminator.train()
-
-            optimizer_D.zero_grad()
-            optimizer_GE.zero_grad()
-
-            # train the model
-            X_, Y_true_pre, Y_fake_pre = self.train_model(X, gaussian_noise,
-                                                          edge_index)
-
-            # get loss
-            loss_D, loss_GE = self.loss_function(X, X_, Y_true_pre, Y_fake_pre,
-                                                 edge_index, criterion)
-
-            loss_D.backward(retain_graph=True)
-            loss_GE.backward()
-
-            optimizer_D.step()
-            optimizer_GE.step()
-
-            # print out log information
-            if self.verbose:
-                score = self.score_function(X, X_, Y_true_pre, Y_fake_pre,
+            epoch_loss_D, epoch_loss_GE = 0, 0
+            for sampled_data in loader:
+                batch_size = sampled_data.batch_size
+                node_idx = sampled_data.node_idx
+                X, edge_index = self.process_graph(sampled_data)
+
+                # Generate noise for constructing fake attribute
+                gaussian_noise = torch.randn(X.shape[0], self.noise_dim).to(
+                    self.device)
+
+                optimizer_D.zero_grad()
+                optimizer_GE.zero_grad()
+
+                # train the model
+                X_, Y_true_pre, Y_fake_pre = self.train_model(
+                    X, gaussian_noise, edge_index)
+
+                edge_index_i = (
+                    edge_index[0] < batch_size).nonzero().to(
+                    self.device)
+                X, X_ = X[:batch_size], X_[:batch_size]
+                Y_true_pre = torch.reshape(Y_true_pre,
+                                           [Y_true_pre.shape[0]])[edge_index_i]
+                Y_fake_pre = torch.reshape(Y_fake_pre,
+                                           [Y_fake_pre.shape[0]])[edge_index_i]
+                edge_index = torch.index_select(
+                    edge_index, 1, torch.reshape(
+                        edge_index_i, [
+                            edge_index_i.shape[0]]))
+
+                # get loss
+                loss_D, loss_GE = self.loss_function(X, X_,
+                                                     Y_true_pre, Y_fake_pre,
+                                                     edge_index, criterion)
+                epoch_loss_D += loss_D.item() * batch_size
+                epoch_loss_GE += loss_GE.item() * batch_size
+
+                loss_D.backward(retain_graph=True)
+                loss_GE.backward()
+
+                optimizer_D.step()
+                optimizer_GE.step()
+
+                score = self.score_function(X, X_,
+                                            Y_true_pre, Y_fake_pre,
                                             edge_index, criterion)
-                print(
-                    "Epoch {:04d}: Loss GE {:.4f} | Loss D {:.4f}"
-                        .format(epoch, loss_GE.item(), loss_D.item()), end='')
-                if y_true is not None:
-                    auc = eval_roc_auc(y_true, score.detach().cpu().numpy())
-                    print(" | AUC {:.4f}".format(auc))
-
-        score = self.score_function(X, X_, Y_true_pre, Y_fake_pre, edge_index,
-                                    criterion)
-        self.decision_scores_ = score.detach().cpu().numpy()
+                scores[node_idx[:batch_size]] = score.detach().cpu().numpy()
+
+                # print out log information
+                if self.verbose:
+                    print(
+                        "Epoch {:04d}: Loss GE {:.4f} | Loss D {:.4f}"
+                        .format(epoch,
+                                epoch_loss_GE / G.x.shape[0],
+                                epoch_loss_D / G.x.shape[0]), end='')
+                    if y_true is not None:
+                        auc = eval_roc_auc(y_true, scores)
+                        print(" | AUC {:.4f}".format(auc))
+
+        self.decision_scores_ = scores
         self._process_decision_scores()
 
         return self
 
     def decision_function(self, G):
         """
-        Description
-        -----------
         Predict raw anomaly score using the fitted detector.
         Outliers are assigned with larger anomaly scores.
 
@@ -235,29 +271,55 @@ def decision_function(self, G):
         check_is_fitted(self, ['generator', 'encoder', 'discriminator'])
 
         # get needed data object from the input data
-        X, edge_index = self.process_graph(G)
+        G.node_idx = torch.arange(G.x.shape[0])
+        G.s = to_dense_adj(G.edge_index)[0]
+        loader = NeighborLoader(G,
+                                [self.num_neigh] * self.num_layers,
+                                batch_size=self.batch_size)
 
         # enable the evaluation mode
         self.generator.eval()
         self.encoder.eval()
         self.discriminator.eval()
 
-        # construct the vector for holding the reconstruction error
-        gaussian_noise = torch.randn(X.shape[0], self.noise_dim).to(
-            self.device)
-        X_, Y_true_pre, Y_fake_pre = self.train_model(X, gaussian_noise,
-                                                      edge_index)
-
         criterion = torch.nn.BCELoss()
-        outlier_scores = self.score_function(X, X_, Y_true_pre, Y_fake_pre,
-                                             edge_index, criterion)
 
-        return outlier_scores.detach().cpu().numpy()
+        scores = np.zeros(G.x.shape[0])
+        for sampled_data in loader:
+            batch_size = sampled_data.batch_size
+            node_idx = sampled_data.node_idx
+
+            X, edge_index = self.process_graph(sampled_data)
+
+            # construct the vector for holding the reconstruction error
+            gaussian_noise = torch.randn(X.shape[0], self.noise_dim).to(
+                self.device)
+            X_, Y_true_pre, Y_fake_pre = self.train_model(X, gaussian_noise,
+                                                          edge_index)
+
+            edge_index_i = (
+                edge_index[0] < batch_size).nonzero().to(
+                self.device)
+            X, X_ = X[:batch_size], X_[:batch_size]
+            Y_true_pre = torch.reshape(Y_true_pre,
+                                       [Y_true_pre.shape[0]])[edge_index_i]
+            Y_fake_pre = torch.reshape(Y_fake_pre,
+                                       [Y_fake_pre.shape[0]])[edge_index_i]
+            edge_index = torch.index_select(
+                edge_index, 1, torch.reshape(
+                    edge_index_i, [
+                        edge_index_i.shape[0]]))
+
+            score = self.score_function(X, X_,
+                                        Y_true_pre, Y_fake_pre,
+                                        edge_index, criterion)
+
+            scores[node_idx[:batch_size]] = score.detach().cpu().numpy()
+
+        return scores
 
     def train_model(self, X, gaussian_noise, edge_index):
         """
-        Description
-        -----------
         Complete the entire process from noise to generator,
         to encoder, and finally to discriminator.
 
@@ -296,8 +358,6 @@ def train_model(self, X, gaussian_noise, edge_index):
     def loss_function(self, X, X_, Y_true_pre, Y_fake_pre, edge_index,
                       criterion):
         """
-        Description
-        -----------
         Obtain the generator and discriminator losses separately.
 
         Parameters
@@ -341,8 +401,6 @@ def loss_function(self, X, X_, Y_true_pre, Y_fake_pre, edge_index,
     def score_function(self, X, X_, Y_true_pre, Y_fake_pre, edge_index,
                        criterion):
         """
-        Description
-        -----------
         Get anomaly score after the model training by weighted context
         reconstruction loss and structure discriminator loss.
 
@@ -387,14 +445,12 @@ def score_function(self, X, X_, Y_true_pre, Y_fake_pre, edge_index,
                 Y_fake_pre_i, Y_fake)
 
         score = self.alpha * attribute_errors + (
-                1 - self.alpha) * structure_errors
+            1 - self.alpha) * structure_errors
 
         return score
 
     def process_graph(self, G):
         """
-        Description
-        -----------
         Process the raw PyG data object into a tuple of sub data
         objects needed for the model.
 
diff --git a/pygod/models/gcnae.py b/pygod/models/gcnae.py
index 3006837..64872b9 100644
--- a/pygod/models/gcnae.py
+++ b/pygod/models/gcnae.py
@@ -5,11 +5,14 @@
 # License: BSD 2 clause
 
 import torch
+import numpy as np
 import torch.nn.functional as F
+from torch_geometric.loader import NeighborLoader
 from sklearn.utils.validation import check_is_fitted
 
 from . import BaseDetector
 from .basic_nn import GCN
+from ..utils.utility import validate_device
 from ..utils.metric import eval_roc_auc
 
 
@@ -42,6 +45,11 @@ class GCNAE(BaseDetector):
         Maximum number of training epoch. Default: ``100``.
     gpu : int
         GPU Index, -1 for using CPU. Default: ``0``.
+    batch_size : int, optional
+        Minibatch size, 0 for full batch training. Default: ``0``.
+    num_neigh : int, optional
+        Number of neighbors in sampling, -1 for all neighbors.
+        Default: ``-1``.
     verbose : bool
         Verbosity mode. Turn on to print out log information.
         Default: ``False``.
@@ -64,6 +72,8 @@ def __init__(self,
                  lr=5e-3,
                  epoch=100,
                  gpu=0,
+                 batch_size=0,
+                 num_neigh=-1,
                  verbose=False):
         super(GCNAE, self).__init__(contamination=contamination)
 
@@ -77,10 +87,9 @@ def __init__(self,
         # training param
         self.lr = lr
         self.epoch = epoch
-        if gpu >= 0 and torch.cuda.is_available():
-            self.device = 'cuda:{}'.format(gpu)
-        else:
-            self.device = 'cpu'
+        self.device = validate_device(gpu)
+        self.batch_size = batch_size
+        self.num_neigh = num_neigh
 
         # other param
         self.verbose = verbose
@@ -88,64 +97,74 @@ def __init__(self,
 
     def fit(self, G, y_true=None):
         """
-        Description
-        -----------
         Fit detector with input data.
 
         Parameters
         ----------
-        G : PyTorch Geometric Data instance (torch_geometric.data.Data)
+        G : torch_geometric.data.Data
             The input data.
-        y_true : numpy.array, optional (default=None)
-            The optional outlier ground truth labels used to monitor the
-            training progress. They are not used to optimize the
-            unsupervised model.
+        y_true : numpy.ndarray, optional
+            The optional outlier ground truth labels used to monitor
+            the training progress. They are not used to optimize the
+            unsupervised model. Default: ``None``.
 
         Returns
         -------
         self : object
             Fitted estimator.
         """
-        x, edge_index = self.process_graph(G)
-
-        self.model = GCN(in_channels=x.shape[1],
+        G.node_idx = torch.arange(G.x.shape[0])
+        if self.batch_size == 0:
+            self.batch_size = G.x.shape[0]
+        loader = NeighborLoader(G,
+                                [self.num_neigh] * self.num_layers,
+                                batch_size=self.batch_size)
+
+        self.model = GCN(in_channels=G.x.shape[1],
                          hidden_channels=self.hid_dim,
                          num_layers=self.num_layers,
-                         out_channels=x.shape[1],
+                         out_channels=G.x.shape[1],
                          dropout=self.dropout,
                          act=self.act).to(self.device)
 
         optimizer = torch.optim.Adam(self.model.parameters(),
                                      lr=self.lr,
                                      weight_decay=self.weight_decay)
-        score = None
-        for epoch in range(self.epoch):
-            self.model.train()
-            x_ = self.model(x, edge_index)
-            loss = F.mse_loss(x_, x)
-
-            optimizer.zero_grad()
-            loss.backward()
-            optimizer.step()
 
-            score = torch.mean(F.mse_loss(x_, x, reduction='none'), dim=1)
+        self.model.train()
+        decision_scores = np.zeros(G.x.shape[0])
+        for epoch in range(self.epoch):
+            epoch_loss = 0
+            for sampled_data in loader:
+                batch_size = sampled_data.batch_size
+                node_idx = sampled_data.node_idx
+                x, edge_index = self.process_graph(sampled_data)
+
+                x_ = self.model(x, edge_index)
+                score = torch.mean(F.mse_loss(x_, x, reduction='none'), dim=1)
+                decision_scores[node_idx[:batch_size]] = score.detach()\
+                                                              .cpu().numpy()
+                loss = torch.mean(score)
+                epoch_loss += loss.item() * batch_size
+
+                optimizer.zero_grad()
+                loss.backward()
+                optimizer.step()
 
             if self.verbose:
                 print("Epoch {:04d}: Loss {:.4f}"
-                      .format(epoch, loss.item()), end='')
+                      .format(epoch, epoch_loss / G.x.shape[0]), end='')
                 if y_true is not None:
-                    auc = eval_roc_auc(y_true, score.detach().cpu().numpy())
+                    auc = eval_roc_auc(y_true, decision_scores)
                     print(" | AUC {:.4f}".format(auc), end='')
                 print()
 
-        self.decision_scores_ = score.detach().cpu().numpy()
+        self.decision_scores_ = decision_scores
         self._process_decision_scores()
         return self
 
     def decision_function(self, G):
         """
-        Description
-        -----------
         Predict raw anomaly score using the fitted detector. Outliers
         are assigned with larger anomaly scores.
 
@@ -160,19 +179,26 @@ def decision_function(self, G):
             The anomaly score of shape :math:`N`.
         """
         check_is_fitted(self, ['model'])
-        self.model.eval()
+        G.node_idx = torch.arange(G.x.shape[0])
+        loader = NeighborLoader(G,
+                                [self.num_neigh] * self.num_layers,
+                                batch_size=self.batch_size)
 
-        x, edge_index = self.process_graph(G)
+        self.model.eval()
+        outlier_scores = np.zeros(G.x.shape[0])
+        for sampled_data in loader:
+            batch_size = sampled_data.batch_size
+            node_idx = sampled_data.node_idx
+            x, edge_index = self.process_graph(sampled_data)
 
-        x_ = self.model(x, edge_index)
-        outlier_scores = torch.mean(F.mse_loss(x_, x, reduction='none'),
-                                    dim=1).detach().cpu().numpy()
+            x_ = self.model(x, edge_index)
+            score = torch.mean(F.mse_loss(x_, x, reduction='none'), dim=1)
+            outlier_scores[node_idx[:batch_size]] = score.detach() \
+                                                         .cpu().numpy()
         return outlier_scores
 
     def process_graph(self, G):
         """
-        Description
-        -----------
         Process the raw PyG data object into a tuple of sub data
         objects needed for the model.
 
diff --git a/pygod/models/guide.py b/pygod/models/guide.py
index 950e1e2..eaaddfe 100644
--- a/pygod/models/guide.py
+++ b/pygod/models/guide.py
@@ -4,18 +4,23 @@
 # Author: Kay Liu <zliu234@uic.edu>
 # License: BSD 2 clause
 
+import os
 import torch
+import hashlib
+import numpy as np
 import networkx as nx
 import torch.nn as nn
 import torch.nn.functional as F
 from torch_geometric.utils import softmax
-from torch_geometric.utils import add_self_loops
 from torch_geometric.nn import MessagePassing
+from torch_geometric.utils import add_self_loops
+from torch_geometric.loader import NeighborLoader
 from networkx.generators.atlas import graph_atlas_g
 from sklearn.utils.validation import check_is_fitted
 
 from . import BaseDetector
 from .basic_nn import GCN
+from ..utils.utility import validate_device
 from ..utils.metric import eval_roc_auc
 
 
@@ -32,15 +37,17 @@ class GUIDE(BaseDetector):
     autoencoders are defined as structure anomaly score and attribute
     anomaly score, respectively.
 
-    Note: The graph preprocesing in the model has high time complexity.
-    It may take longer than you expect.
+    Note: The calculation of node motif degree in preprocesing has high
+    time complexity. It may take longer than you expect.
 
     See :cite:`yuan2021higher` for details.
 
     Parameters
     ----------
-    hid_dim :  int, optional
-        Hidden dimension of model. Default: ``0``.
+    a_hid :  int, optional
+        Hidden dimension for attribute autoencoder. Default: ``32``.
+    s_hid :  int, optional
+        Hidden dimension for structure autoencoder. Default: ``4``.
     num_layers : int, optional
         Total number of layers in autoencoders. Default: ``4``.
     dropout : float, optional
@@ -63,12 +70,20 @@ class GUIDE(BaseDetector):
         Maximum number of training epoch. Default: ``10``.
     gpu : int
         GPU Index, -1 for using CPU. Default: ``0``.
+    batch_size : int, optional
+        Minibatch size, 0 for full batch training. Default: ``0``.
+    num_neigh : int, optional
+        Number of neighbors in sampling, -1 for all neighbors.
+        Default: ``-1``.
     graphlet_size : int
         The maximum graphlet size used to compute structure input.
         Default: ``4``.
     selected_motif : bool
         Use selected motifs which are defined in the original paper.
-        Default: ``False``.
+        Default: ``True``.
+    cache_dir : str
+        The directory for the node motif degree caching.
+        Default: ``None``.
     verbose : bool
         Verbosity mode. Turn on to print out log information.
         Default: ``False``.
@@ -82,7 +97,8 @@ class GUIDE(BaseDetector):
     """
 
     def __init__(self,
-                 hid_dim=64,
+                 a_hid=32,
+                 s_hid=4,
                  num_layers=4,
                  dropout=0.3,
                  weight_decay=0.,
@@ -92,13 +108,17 @@ def __init__(self,
                  lr=0.001,
                  epoch=10,
                  gpu=0,
+                 batch_size=0,
+                 num_neigh=-1,
                  graphlet_size=4,
-                 selected_motif=False,
+                 selected_motif=True,
+                 cache_dir=None,
                  verbose=False):
         super(GUIDE, self).__init__(contamination=contamination)
 
         # model param
-        self.hid_dim = hid_dim
+        self.a_hid = a_hid
+        self.s_hid = s_hid
         self.num_layers = num_layers
         self.dropout = dropout
         self.weight_decay = weight_decay
@@ -108,10 +128,9 @@ def __init__(self,
         # training param
         self.lr = lr
         self.epoch = epoch
-        if gpu >= 0 and torch.cuda.is_available():
-            self.device = 'cuda:{}'.format(gpu)
-        else:
-            self.device = 'cpu'
+        self.device = validate_device(gpu)
+        self.batch_size = batch_size
+        self.num_neigh = num_neigh
 
         # other param
         self.graphlet_size = graphlet_size
@@ -122,32 +141,40 @@ def __init__(self,
         self.verbose = verbose
         self.model = None
 
+        if cache_dir is None:
+            cache_dir = os.path.join(os.path.expanduser('~'), '.pygod')
+        self.cache_dir = cache_dir
+
     def fit(self, G, y_true=None):
         """
-        Description
-        -----------
         Fit detector with input data.
 
         Parameters
         ----------
-        G : PyTorch Geometric Data instance (torch_geometric.data.Data)
+        G : torch_geometric.data.Data
             The input data.
-        y_true : numpy.array, optional (default=None)
-            The optional outlier ground truth labels used to monitor the
-            training progress. They are not used to optimize the
-            unsupervised model.
+        y_true : numpy.ndarray, optional
+            The optional outlier ground truth labels used to monitor
+            the training progress. They are not used to optimize the
+            unsupervised model. Default: ``None``.
 
         Returns
         -------
         self : object
             Fitted estimator.
         """
-
-        x, s, edge_index = self.process_graph(G)
-
-        self.model = GUIDE_Base(x_dim=x.shape[1],
-                                s_dim=s.shape[1],
-                                hid_dim=self.hid_dim,
+        G.node_idx = torch.arange(G.x.shape[0])
+        G.s = self._get_nmf(G, self.cache_dir)
+        if self.batch_size == 0:
+            self.batch_size = G.x.shape[0]
+        loader = NeighborLoader(G,
+                                [self.num_neigh] * self.num_layers,
+                                batch_size=self.batch_size)
+
+        self.model = GUIDE_Base(a_dim=G.x.shape[1],
+                                s_dim=G.s.shape[1],
+                                a_hid=self.a_hid,
+                                s_hid=self.s_hid,
                                 num_layers=self.num_layers,
                                 dropout=self.dropout,
                                 act=self.act).to(self.device)
@@ -156,33 +183,43 @@ def fit(self, G, y_true=None):
                                      lr=self.lr,
                                      weight_decay=self.weight_decay)
 
-        score = None
+        self.model.train()
+        decision_scores = np.zeros(G.x.shape[0])
         for epoch in range(self.epoch):
-            self.model.train()
-            x_, s_ = self.model(x, s, edge_index)
-            score = self.loss_func(x, x_, s, s_)
-            loss = torch.mean(score)
-
-            optimizer.zero_grad()
-            loss.backward()
-            optimizer.step()
+            epoch_loss = 0
+            for sampled_data in loader:
+                batch_size = sampled_data.batch_size
+                node_idx = sampled_data.node_idx
+                x, s, edge_index = self.process_graph(G)
+
+                x_, s_ = self.model(x, s, edge_index)
+                score = self.loss_func(x[:batch_size],
+                                       x_[:batch_size],
+                                       s[:batch_size],
+                                       s_[:batch_size])
+                decision_scores[node_idx[:batch_size]] = score.detach() \
+                                                              .cpu().numpy()
+                loss = torch.mean(score)
+                epoch_loss += loss.item() * batch_size
+
+                optimizer.zero_grad()
+                loss.backward()
+                optimizer.step()
 
             if self.verbose:
                 print("Epoch {:04d}: Loss {:.4f}"
-                      .format(epoch, loss.item()), end='')
+                      .format(epoch, epoch_loss / G.x.shape[0]), end='')
                 if y_true is not None:
-                    auc = eval_roc_auc(y_true, score.detach().cpu().numpy())
+                    auc = eval_roc_auc(y_true, decision_scores)
                     print(" | AUC {:.4f}".format(auc), end='')
                 print()
 
-        self.decision_scores_ = score.detach().cpu().numpy()
+        self.decision_scores_ = decision_scores
         self._process_decision_scores()
         return self
 
     def decision_function(self, G):
         """
-        Description
-        -----------
         Predict raw anomaly score using the fitted detector. Outliers
         are assigned with larger anomaly scores.
 
@@ -197,25 +234,35 @@ def decision_function(self, G):
             The anomaly score of shape :math:`N`.
         """
         check_is_fitted(self, ['model'])
+        G.node_idx = torch.arange(G.x.shape[0])
+        G.s = self._get_nmf(G, self.cache_dir)
+        loader = NeighborLoader(G,
+                                [self.num_neigh] * self.num_layers,
+                                batch_size=self.batch_size)
 
-        # get needed data object from the input data
-        x, s, edge_index = self.process_graph(G)
-
-        # enable the evaluation mode
         self.model.eval()
+        outlier_scores = np.zeros(G.x.shape[0])
+        for sampled_data in loader:
+            batch_size = sampled_data.batch_size
+            node_idx = sampled_data.node_idx
 
-        # construct the vector for holding the reconstruction error
-        x_, s_ = self.model(x, s, edge_index)
-        outlier_scores = self.loss_func(x, x_, s, s_)
-        return outlier_scores.detach().cpu().numpy()
+            x, s, edge_index = self.process_graph(sampled_data)
+
+            x_, s_ = self.model(x, s, edge_index)
+            score = self.loss_func(x[:batch_size],
+                                   x_[:batch_size],
+                                   s[:batch_size],
+                                   s_[:batch_size])
+
+            outlier_scores[node_idx[:batch_size]] = score.detach()\
+                                                         .cpu().numpy()
+
+        return outlier_scores
 
     def process_graph(self, G):
         """
-        Description
-        -----------
         Process the raw PyG data object into a tuple of sub data
-        objects needed for the model. Part of this function is adapted
-        from https://github.com/benedekrozemberczki/OrbitalFeatures.
+        objects needed for the model.
 
         Parameters
         ----------
@@ -231,90 +278,8 @@ def process_graph(self, G):
         edge_index : torch.Tensor
             Edge list of the graph.
         """
-        edge_index = G.edge_index
-        g = nx.from_edgelist(edge_index.T.tolist())
-
-        # create edge subsets
-        edge_subsets = dict()
-        subsets = [[edge[0], edge[1]] for edge in g.edges()]
-        edge_subsets[2] = subsets
-        unique_subsets = dict()
-        for i in range(3, self.graphlet_size + 1):
-            for subset in subsets:
-                for node in subset:
-                    for neb in g.neighbors(node):
-                        new_subset = subset + [neb]
-                        if len(set(new_subset)) == i:
-                            new_subset.sort()
-                            unique_subsets[tuple(new_subset)] = 1
-            subsets = [list(k) for k, v in unique_subsets.items()]
-            edge_subsets[i] = subsets
-            unique_subsets = dict()
-
-        # enumerate graphs
-        graphs = graph_atlas_g()
-        interesting_graphs = {i: [] for i in
-                              range(2, self.graphlet_size + 1)}
-        for graph in graphs:
-            if 1 < graph.number_of_nodes() < self.graphlet_size + 1:
-                if nx.is_connected(graph):
-                    interesting_graphs[graph.number_of_nodes()].append(
-                        graph)
-
-        # enumerate categories
-        main_index = 0
-        categories = dict()
-        for size, graphs in interesting_graphs.items():
-            categories[size] = dict()
-            for index, graph in enumerate(graphs):
-                categories[size][index] = dict()
-                degrees = list(
-                    set([graph.degree(node) for node in graph.nodes()]))
-                for degree in degrees:
-                    categories[size][index][degree] = main_index
-                    main_index += 1
-        unique_motif_count = main_index
-
-        # setup feature
-        features = {node: {i: 0 for i in range(unique_motif_count)}
-                    for node in g.nodes()}
-        for size, node_lists in edge_subsets.items():
-            graphs = interesting_graphs[size]
-            for nodes in node_lists:
-                sub_gr = g.subgraph(nodes)
-                for index, graph in enumerate(graphs):
-                    if nx.is_isomorphic(sub_gr, graph):
-                        for node in sub_gr.nodes():
-                            features[node][categories[size][index][
-                                sub_gr.degree(node)]] += 1
-                        break
-
-        motifs = [[n] + [features[n][i] for i in range(
-            unique_motif_count)] for n in g.nodes()]
-        motifs = torch.Tensor(motifs)
-        motifs = motifs[torch.sort(motifs[:, 0]).indices, 1:]
-
-        if self.selected_motif:
-            # use motif selected in the original paper only
-            s = torch.zeros((G.x.shape[0], 6))
-            # m31
-            s[:, 0] = motifs[:, 3]
-            # m32
-            s[:, 1] = motifs[:, 1] + motifs[:, 2]
-            # m41
-            s[:, 2] = motifs[:, 14]
-            # m42
-            s[:, 3] = motifs[:, 12] + motifs[:, 13]
-            # m43
-            s[:, 4] = motifs[:, 11]
-            # node degree
-            s[:, 5] = motifs[:, 0]
-        else:
-            # use graphlet degree
-            s = motifs
-
-        edge_index = edge_index.to(self.device)
-        s = s.to(self.device)
+        edge_index = G.edge_index.to(self.device)
+        s = G.s.to(self.device)
         x = G.x.to(self.device)
 
         return x, s, edge_index
@@ -332,26 +297,142 @@ def loss_func(self, x, x_, s, s_):
                 structure_errors
         return score
 
+    def _get_nmf(self, G, cache_dir):
+        """
+        Calculation of Node Motif Degree / Graphlet Degree
+        Distribution. Part of this function is adapted
+        from https://github.com/benedekrozemberczki/OrbitalFeatures.
+
+        Parameters
+        ----------
+        G : PyTorch Geometric Data instance (torch_geometric.data.Data)
+            The input data.
+        cache_dir : str
+            The directory for the node motif degree caching
+
+        Returns
+        -------
+        s : torch.Tensor
+            Structure matrix (node motif degree/graphlet degree)
+        """
+        hash_func = hashlib.sha1()
+        hash_func.update(str(G).encode('utf-8'))
+        file_name = 'nmd_' + str(hash_func.hexdigest()[:8]) + \
+                    str(self.graphlet_size) + \
+                    str(self.selected_motif)[0] + '.pt'
+        file_path = os.path.join(cache_dir, file_name)
+        if os.path.exists(file_path):
+            s = torch.load(file_path)
+        else:
+            edge_index = G.edge_index
+            g = nx.from_edgelist(edge_index.T.tolist())
+
+            # create edge subsets
+            edge_subsets = dict()
+            subsets = [[edge[0], edge[1]] for edge in g.edges()]
+            edge_subsets[2] = subsets
+            unique_subsets = dict()
+            for i in range(3, self.graphlet_size + 1):
+                for subset in subsets:
+                    for node in subset:
+                        for neb in g.neighbors(node):
+                            new_subset = subset + [neb]
+                            if len(set(new_subset)) == i:
+                                new_subset.sort()
+                                unique_subsets[tuple(new_subset)] = 1
+                subsets = [list(k) for k, v in unique_subsets.items()]
+                edge_subsets[i] = subsets
+                unique_subsets = dict()
+
+            # enumerate graphs
+            graphs = graph_atlas_g()
+            interesting_graphs = {i: [] for i in
+                                  range(2, self.graphlet_size + 1)}
+            for graph in graphs:
+                if 1 < graph.number_of_nodes() < self.graphlet_size + 1:
+                    if nx.is_connected(graph):
+                        interesting_graphs[graph.number_of_nodes()].append(
+                            graph)
+
+            # enumerate categories
+            main_index = 0
+            categories = dict()
+            for size, graphs in interesting_graphs.items():
+                categories[size] = dict()
+                for index, graph in enumerate(graphs):
+                    categories[size][index] = dict()
+                    degrees = list(
+                        set([graph.degree(node) for node in graph.nodes()]))
+                    for degree in degrees:
+                        categories[size][index][degree] = main_index
+                        main_index += 1
+            unique_motif_count = main_index
+
+            # setup feature
+            features = {node: {i: 0 for i in range(unique_motif_count)}
+                        for node in g.nodes()}
+            for size, node_lists in edge_subsets.items():
+                graphs = interesting_graphs[size]
+                for nodes in node_lists:
+                    sub_gr = g.subgraph(nodes)
+                    for index, graph in enumerate(graphs):
+                        if nx.is_isomorphic(sub_gr, graph):
+                            for node in sub_gr.nodes():
+                                features[node][categories[size][index][
+                                    sub_gr.degree(node)]] += 1
+                            break
+
+            motifs = [[n] + [features[n][i] for i in range(
+                unique_motif_count)] for n in g.nodes()]
+            motifs = torch.Tensor(motifs)
+            motifs = motifs[torch.sort(motifs[:, 0]).indices, 1:]
+
+            if self.selected_motif:
+                # use motif selected in the original paper only
+                s = torch.zeros((G.x.shape[0], 6))
+                # m31
+                s[:, 0] = motifs[:, 3]
+                # m32
+                s[:, 1] = motifs[:, 1] + motifs[:, 2]
+                # m41
+                s[:, 2] = motifs[:, 14]
+                # m42
+                s[:, 3] = motifs[:, 12] + motifs[:, 13]
+                # m43
+                s[:, 4] = motifs[:, 11]
+                # node degree
+                s[:, 5] = motifs[:, 0]
+            else:
+                # use graphlet degree
+                s = motifs
+
+            if not os.path.exists(cache_dir):
+                os.makedirs(cache_dir)
+            torch.save(s, file_path)
+
+        return s
+
 
 class GUIDE_Base(nn.Module):
     def __init__(self,
-                 x_dim,
+                 a_dim,
                  s_dim,
-                 hid_dim,
+                 a_hid,
+                 s_hid,
                  num_layers,
                  dropout,
                  act):
         super(GUIDE_Base, self).__init__()
 
-        self.attr_ae = GCN(in_channels=x_dim,
-                           hidden_channels=hid_dim,
+        self.attr_ae = GCN(in_channels=a_dim,
+                           hidden_channels=a_hid,
                            num_layers=num_layers,
-                           out_channels=x_dim,
+                           out_channels=a_dim,
                            dropout=dropout,
                            act=act)
 
         self.struct_ae = GNA(in_channels=s_dim,
-                             hidden_channels=hid_dim,
+                             hidden_channels=s_hid,
                              num_layers=num_layers,
                              out_channels=s_dim,
                              dropout=dropout,
diff --git a/pygod/models/mlpae.py b/pygod/models/mlpae.py
index 22b0e4b..352bbe1 100644
--- a/pygod/models/mlpae.py
+++ b/pygod/models/mlpae.py
@@ -5,12 +5,16 @@
 # License: BSD 2 clause
 
 import torch
+import numpy as np
 import torch.nn.functional as F
+from torch.utils.data import DataLoader
 from sklearn.utils.validation import check_is_fitted
 
 from . import BaseDetector
 from .basic_nn import MLP
+from ..utils.utility import validate_device
 from ..utils.metric import eval_roc_auc
+from ..utils.dataset import PlainDataset
 
 
 class MLPAE(BaseDetector):
@@ -42,6 +46,8 @@ class MLPAE(BaseDetector):
         Maximum number of training epoch. Default: ``5``.
     gpu : int
         GPU Index, -1 for using CPU. Default: ``0``.
+    batch_size : int, optional
+        Minibatch size, 0 for full batch training. Default: ``0``.
     verbose : bool
         Verbosity mode. Turn on to print out log information.
         Default: ``False``.
@@ -63,6 +69,7 @@ def __init__(self,
                  lr=5e-3,
                  epoch=5,
                  gpu=0,
+                 batch_size=0,
                  verbose=False):
         super(MLPAE, self).__init__(contamination=contamination)
 
@@ -76,10 +83,8 @@ def __init__(self,
         # training param
         self.lr = lr
         self.epoch = epoch
-        if gpu >= 0 and torch.cuda.is_available():
-            self.device = 'cuda:{}'.format(gpu)
-        else:
-            self.device = 'cpu'
+        self.device = validate_device(gpu)
+        self.batch_size = batch_size
 
         # other param
         self.verbose = verbose
@@ -87,29 +92,31 @@ def __init__(self,
 
     def fit(self, G, y_true=None):
         """
-        Description
-        -----------
         Fit detector with input data.
 
         Parameters
         ----------
-        G : PyTorch Geometric Data instance (torch_geometric.data.Data)
+        G : torch_geometric.data.Data
             The input data.
-        y_true : numpy.array, optional (default=None)
-            The optional outlier ground truth labels used to monitor the
-            training progress. They are not used to optimize the
-            unsupervised model.
+        y_true : numpy.ndarray, optional
+            The optional outlier ground truth labels used to monitor
+            the training progress. They are not used to optimize the
+            unsupervised model. Default: ``None``.
 
         Returns
         -------
         self : object
             Fitted estimator.
         """
-        x = self.process_graph(G)
+        full_x = self.process_graph(G)
+        dataset = PlainDataset(full_x)
+        if self.batch_size == 0:
+            self.batch_size = G.x.shape[0]
+        loader = DataLoader(dataset, batch_size=self.batch_size)
 
-        self.model = MLP(in_channels=x.shape[1],
+        self.model = MLP(in_channels=G.x.shape[1],
                          hidden_channels=self.hid_dim,
-                         out_channels=x.shape[1],
+                         out_channels=G.x.shape[1],
                          num_layers=self.num_layers,
                          dropout=self.dropout,
                          act=self.act).to(self.device)
@@ -118,34 +125,35 @@ def fit(self, G, y_true=None):
                                      lr=self.lr,
                                      weight_decay=self.weight_decay)
 
-        score = None
+        self.model.train()
+        decision_scores = np.zeros(full_x.shape[0])
         for epoch in range(self.epoch):
-            self.model.train()
-            x_ = self.model(x)
-            loss = F.mse_loss(x_, x)
-
-            optimizer.zero_grad()
-            loss.backward()
-            optimizer.step()
-
-            score = torch.mean(F.mse_loss(x_, x, reduction='none'), dim=1)
+            epoch_loss = 0
+            for x, node_idx in loader:
+                x_ = self.model(x)
+                score = torch.mean(F.mse_loss(x_, x, reduction='none'), dim=1)
+                decision_scores[node_idx] = score.detach().cpu().numpy()
+                loss = torch.mean(score)
+                epoch_loss += loss.item() * x.shape[0]
+
+                optimizer.zero_grad()
+                loss.backward()
+                optimizer.step()
 
             if self.verbose:
                 print("Epoch {:04d}: Loss {:.4f}"
-                      .format(epoch, loss.item()), end='')
+                      .format(epoch, epoch_loss / G.x.shape[0]), end='')
                 if y_true is not None:
-                    auc = eval_roc_auc(y_true, score.detach().cpu().numpy())
+                    auc = eval_roc_auc(y_true, decision_scores)
                     print(" | AUC {:.4f}".format(auc), end='')
                 print()
 
-        self.decision_scores_ = score.detach().cpu().numpy()
+        self.decision_scores_ = decision_scores
         self._process_decision_scores()
         return self
 
     def decision_function(self, G):
         """
-        Description
-        -----------
         Predict raw anomaly score using the fitted detector. Outliers
         are assigned with larger anomaly scores.
 
@@ -160,20 +168,20 @@ def decision_function(self, G):
             The anomaly score of shape :math:`N`.
         """
         check_is_fitted(self, ['model'])
-        self.model.eval()
+        full_x = self.process_graph(G)
+        dataset = PlainDataset(full_x)
+        loader = DataLoader(dataset, batch_size=self.batch_size)
 
-        x = self.process_graph(G)
-        x = x.to(self.device)
-
-        x_ = self.model(x)
-        outlier_scores = torch.mean(F.mse_loss(x_, x, reduction='none')
-                                    , dim=1).detach().cpu().numpy()
+        self.model.eval()
+        outlier_scores = np.zeros(full_x.shape[0])
+        for x, node_idx in loader:
+            x_ = self.model(x)
+            score = torch.mean(F.mse_loss(x_, x, reduction='none'), dim=1)
+            outlier_scores[node_idx] = score.detach().cpu().numpy()
         return outlier_scores
 
     def process_graph(self, G):
         """
-        Description
-        -----------
         Process the raw PyG data object into a tuple of sub data
         objects needed for the model.
 
@@ -188,5 +196,4 @@ def process_graph(self, G):
             Attribute (feature) of nodes.
         """
         x = G.x.to(self.device)
-
         return x
diff --git a/pygod/models/ocgnn.py b/pygod/models/ocgnn.py
index 2720b37..e139e16 100644
--- a/pygod/models/ocgnn.py
+++ b/pygod/models/ocgnn.py
@@ -13,7 +13,9 @@
 from torch_sparse import SparseTensor
 
 from . import BaseDetector
+from ..utils.utility import validate_device
 from ..utils.metric import eval_roc_auc
+from torch_geometric.loader import NeighborLoader
 
 
 class GCN_base(nn.Module):
@@ -47,7 +49,7 @@ class OCGNN(BaseDetector):
     Networks): OCGNN is an anomaly detector that measures the distance of
     anomaly to the centroid, in the similar fashion to the support vector
     machine, but in the embedding space after feeding towards several layers
-     of GCN.
+    of GCN.
 
     See :cite:`wang2021one` for details.
 
@@ -85,6 +87,11 @@ class OCGNN(BaseDetector):
     verbose : bool
         Verbosity mode. Turn on to print out log information.
         Defaults: ``False``.
+    batch_size : int, optional
+        Minibatch size, 0 for full batch training. Default: ``0``.
+    num_neigh : int, optional
+        Number of neighbors in sampling, -1 for all neighbors.
+        Default: ``-1``.
 
     Examples
     --------
@@ -107,7 +114,9 @@ def __init__(self,
                  epoch=5,
                  warmup_epoch=2,
                  verbose=False,
-                 act=F.relu):
+                 act=F.relu,
+                 batch_size = 0,
+                 num_neigh = -1):
         super(OCGNN, self).__init__(contamination=contamination)
         self.dropout = dropout
         self.n_hidden = n_hidden
@@ -121,11 +130,9 @@ def __init__(self,
         self.epoch = epoch
         self.warmup_epoch = warmup_epoch
         self.act = act
-
-        if gpu >= 0 and torch.cuda.is_available():
-            self.device = 'cuda:{}'.format(gpu)
-        else:
-            self.device = 'cpu'
+        self.device = validate_device(gpu)
+        self.batch_size = batch_size
+        self.num_neigh = num_neigh
 
         # other param
         self.verbose = verbose
@@ -133,8 +140,6 @@ def __init__(self,
 
     def init_center(self, x, edge_index):
         """
-        Descriptions
-        ----------
         Initialize hypersphere center c as the mean from
         an initial forward pass on the data.
   
@@ -167,8 +172,6 @@ def init_center(self, x, edge_index):
 
     def get_radius(self, dist):
         """
-        Description
-        ----------
         Optimally solve for radius R via the (1-nu)-quantile of distances.
         
         Parameters
@@ -187,8 +190,6 @@ def get_radius(self, dist):
 
     def anomaly_scores(self, outputs):
         """
-        Description
-        ----------
         Calculate the anomaly score given by Euclidean distance to the center.
         
         Parameters
@@ -209,8 +210,6 @@ def anomaly_scores(self, outputs):
 
     def loss_function(self, outputs, update=False):
         """
-        Description
-        ----------
         Calculate the loss in paper Equation (4)
         
         Parameters
@@ -239,26 +238,30 @@ def loss_function(self, outputs, update=False):
 
     def fit(self, G, y_true=None):
         """
-        Description
-        -----------
         Fit detector with input data.
 
         Parameters
         ----------
-        G : PyTorch Geometric Data instance (torch_geometric.data.Data)
+        G : torch_geometric.data.Data
             The input data.
-        y_true : numpy.array, optional (default=None)
-            The optional outlier ground truth labels used to monitor the
-            training progress. They are not used to optimize the
-            unsupervised model.
+        y_true : numpy.ndarray, optional
+            The optional outlier ground truth labels used to monitor
+            the training progress. They are not used to optimize the
+            unsupervised model. Default: ``None``.
 
         Returns
         -------
         self : object
             Fitted estimator.
         """
-        x, adj, edge_index = self.process_graph(G)
-        self.in_feats = x.shape[1]
+        G.node_idx = torch.arange(G.x.shape[0])
+        if self.batch_size == 0:
+            self.batch_size = G.x.shape[0]
+        loader = NeighborLoader(G,
+                                [self.num_neigh] * self.n_layers,
+                                batch_size=self.batch_size)
+ 
+        self.in_feats = G.x.shape[1]
 
         # initialize the model and optimizer
         self.model = GCN_base(self.in_feats,
@@ -277,35 +280,41 @@ def fit(self, G, y_true=None):
         self.model = self.model.to(self.device)
         # training the model
         self.model.train()
-
-        score = None
+        decision_scores = np.zeros(G.x.shape[0])
         for cur_epoch in range(self.epoch):
-            self.model.zero_grad()
-            outputs = self.model(x, edge_index)
-            loss, dist, score = self.loss_function(outputs)
-            if self.warmup_epoch is not None and cur_epoch < self.warmup_epoch:
-                self.data_center = self.init_center(x, edge_index)
-                self.radius = torch.tensor(self.get_radius(dist),
+            epoch_loss = 0.0
+            for sampled_data in loader:
+                batch_size = sampled_data.batch_size
+                node_idx = sampled_data.node_idx
+                x, adj, edge_index = self.process_graph(sampled_data)
+                outputs = self.model(x, edge_index)
+                loss, dist, score = self.loss_function(outputs)
+                epoch_loss += loss.item() * batch_size
+                if self.warmup_epoch is not None and cur_epoch < self.warmup_epoch:
+                    self.data_center = self.init_center(x, edge_index)
+                    self.radius = torch.tensor(self.get_radius(dist),
                                            device=self.device)
-            loss.backward()
-            self.optimizer.step()
+                
+                decision_scores[node_idx[:batch_size]] = score.detach()\
+                                                              .cpu().numpy()
+                self.model.zero_grad()
+                loss.backward()
+                self.optimizer.step()
 
             if self.verbose:
                 print("Epoch {:04d}: Loss {:.4f}"
-                      .format(cur_epoch, loss.item()), end='')
+                      .format(cur_epoch,  epoch_loss / G.x.shape[0]), end='')
                 if y_true is not None:
-                    auc = eval_roc_auc(y_true, score.detach().cpu().numpy())
+                    auc = eval_roc_auc(y_true, decision_scores)
                     print(" | AUC {:.4f}".format(auc), end='')
                 print()
 
-        self.decision_scores_ = score.detach().cpu().numpy()
+        self.decision_scores_ = decision_scores
         self._process_decision_scores()
         return self
 
     def process_graph(self, G):
         """
-        Description
-        -----------
         Process the raw PyG data object into a tuple of sub data
         objects needed for the model.
 
@@ -359,15 +368,29 @@ def decision_function(self, G):
         """
         check_is_fitted(self, ['model'])
 
-        # get needed data object from the input data
-        x, adj, edge_index = self.process_graph(G)
+       # get needed data object from the input data
+        G.node_idx = torch.arange(G.x.shape[0])
+        loader = NeighborLoader(G,
+                                [self.num_neigh] * self.n_layers,
+                                batch_size=self.batch_size)
+
 
         # enable the evaluation mode
         self.model.eval()
 
         # construct the vector for holding the reconstruction error
         # outlier_scores = torch.zeros([attrs.shape[0], ])
-        outputs = self.model(x, edge_index)
-        loss, dist, score = self.loss_function(outputs)
+        outlier_scores = np.zeros(G.x.shape[0])
+        for sampled_data in loader:
+            batch_size = sampled_data.batch_size
+            node_idx = sampled_data.node_idx
+
+            x, adj, edge_index = self.process_graph(sampled_data)
+
+            outputs = self.model(x, edge_index)
+            loss, dist, score = self.loss_function(outputs)
+
+            outlier_scores[node_idx[:batch_size]] = score.detach() \
+                .cpu().numpy()
 
-        return score.detach().cpu().numpy()
+        return outlier_scores
diff --git a/pygod/models/one.py b/pygod/models/one.py
index 7255ec5..8cd8cfd 100644
--- a/pygod/models/one.py
+++ b/pygod/models/one.py
@@ -90,18 +90,16 @@ def __init__(self,
 
     def fit(self, G, y_true=None):
         """
-        Description
-        -----------
         Fit detector with input data.
 
         Parameters
         ----------
-        G : PyTorch Geometric Data instance (torch_geometric.data.Data)
+        G : torch_geometric.data.Data
             The input data.
-        y_true : numpy.array, optional (default=None)
-            The optional outlier ground truth labels used to monitor the
-            training progress. They are not used to optimize the
-            unsupervised model.
+        y_true : numpy.ndarray, optional
+            The optional outlier ground truth labels used to monitor
+            the training progress. They are not used to optimize the
+            unsupervised model. Default: ``None``.
 
         Returns
         -------
@@ -350,8 +348,6 @@ def fit(self, G, y_true=None):
 
     def decision_function(self, G):
         """
-        Description
-        -----------
         Predict raw anomaly score using the fitted detector. Outliers
         are assigned with larger anomaly scores.
 
@@ -376,8 +372,6 @@ def decision_function(self, G):
 
     def process_graph(self, G):
         """
-        Description
-        -----------
         Process the raw PyG data object into a tuple of sub data
         objects needed for the model.
 
@@ -420,8 +414,6 @@ def calc_lossValues(self,
                         beta,
                         gamma):
         """
-        Description
-        -----------
         Calculate the loss. This function is called inside the fit() function.
 
         Parameters
@@ -461,8 +453,6 @@ def cal_outlierScore(self,
                          A,
                          C):
         """
-        Description
-        -----------
         Calculate the outlier scores.
 
         Parameters
diff --git a/pygod/utils/dataset.py b/pygod/utils/dataset.py
new file mode 100644
index 0000000..617d6b7
--- /dev/null
+++ b/pygod/utils/dataset.py
@@ -0,0 +1,12 @@
+from torch.utils.data import Dataset
+
+
+class PlainDataset(Dataset):
+    def __init__(self, x):
+        self.x = x
+
+    def __len__(self):
+        return len(self.x)
+
+    def __getitem__(self, idx):
+        return self.x[idx], idx
diff --git a/pygod/utils/metric.py b/pygod/utils/metric.py
index bc4efeb..ee504e3 100644
--- a/pygod/utils/metric.py
+++ b/pygod/utils/metric.py
@@ -5,20 +5,18 @@
 # Author: Yingtong Dou <ytongdou@gmail.com>
 # License: BSD 2 clause
 
-from sklearn.metrics import roc_auc_score
+from sklearn.metrics import roc_auc_score, average_precision_score
 
 
 def eval_roc_auc(labels, pred):
     """
-    Description
-    -----------
     ROC-AUC score for binary classification.
 
     Parameters
     ----------
     labels : numpy.array
         Labels in shape of ``(N, )``, where 1 represents outliers,
-        0 represents normal nodes.
+        0 represents normal instances.
     pred : numpy.array
         Outlier scores in shape of ``(N, )``.
 
@@ -35,15 +33,13 @@ def eval_roc_auc(labels, pred):
 
 def eval_recall_at_k(labels, pred, k, threshold=0.5):
     """
-    Description
-    -----------
     Recall score for top k instances with the highest outlier scores.
 
     Parameters
     ----------
     labels : numpy.array
         Labels in shape of ``(N, )``, where 1 represents outliers,
-        0 represents normal nodes.
+        0 represents normal instances.
     pred : numpy.array
         Outlier scores in shape of ``(N, )``.
     k : int
@@ -73,15 +69,13 @@ def eval_recall_at_k(labels, pred, k, threshold=0.5):
 
 def eval_precision_at_k(labels, pred, k, threshold=0.5):
     """
-    Description
-    -----------
     Precision score for top k instances with the highest outlier scores.
 
     Parameters
     ----------
     labels : numpy.array
         Labels in shape of ``(N, )``, where 1 represents outliers,
-        0 represents normal nodes.
+        0 represents normal instances.
     pred : numpy.array
         Outlier scores in shape of ``(N, )``.
     k : int
@@ -107,3 +101,26 @@ def eval_precision_at_k(labels, pred, k, threshold=0.5):
     precision_at_k = n_true_and_pred_k / n_pred_k if n_pred_k != 0 else 0.0
 
     return precision_at_k
+
+
+def eval_average_precision(labels, pred):
+    """
+    Average precision score for binary classification.
+
+    Parameters
+    ----------
+    labels : numpy.array
+        Labels in shape of ``(N, )``, where 1 represents outliers,
+        0 represents normal instances.
+    pred : numpy.array
+        Outlier scores in shape of ``(N, )``.
+
+    Returns
+    -------
+    ap : float
+        Average precision score.
+    """
+
+    # outlier detection is a binary classification problem
+    ap = average_precision_score(y_true=labels, y_score=pred)
+    return ap
diff --git a/pygod/utils/outlier_generator.py b/pygod/utils/outlier_generator.py
index ae96851..c364419 100644
--- a/pygod/utils/outlier_generator.py
+++ b/pygod/utils/outlier_generator.py
@@ -5,60 +5,77 @@
 """
 # Author: Yingtong Dou <ytongdou@gmail.com>
 # License: BSD 2 clause
+import math
+import random
 
 import numpy as np
 import torch
+from torch_geometric.data import Data
 
-# set the random seed
-np.random.seed(999)
+from .utility import check_parameter
 
 
-def gen_structure_outliers(data, m, n):
+def gen_structure_outliers(data, m, n, random_state=None):
     """Generating structural outliers according to paper
     "Deep Anomaly Detection on Attributed Networks"
     <https://epubs.siam.org/doi/abs/10.1137/1.9781611975673.67>.
-
     We randomly select m nodes from the network and then make those nodes
     fully connected, and then all the m nodes in the clique are regarded as
     outliers. We iteratively repeat this process until a number of n
     cliques are generated and the total number of structural outliers is m×n.
-
     Parameters
     ----------
-    data : torch_geometric.datasets.Data
-        PyG dataset object.
+    data : PyTorch Geometric Data instance (torch_geometric.data.Data)
+        The input data.
     m : int
         Number nodes in the outlier cliques.
     n : int
         Number of outlier cliques.
+    random_state : int, optional
+        The seed to control the randomness, Default: ``None``.
+
     Returns
     -------
-    data : torch_geometric.datasets.Data
+    data : PyTorch Geometric Data instance (torch_geometric.data.Data)
         The structural outlier graph with injected edges.
     y_outlier : torch.Tensor
         The outlier label tensor where 1 represents outliers and 0 represents
         regular nodes.
     """
 
-    node_set = set(range(data.x.shape[0]))
-    new_edges = []
-    sample_indices = []
-    for i in range(0, n):
+    if not isinstance(data, Data):
+        raise TypeError("data should be torch_geometric.data.Data")
 
-        sampled_idx = np.random.choice(list(node_set), size=m, replace=False)
+    if isinstance(m, int):
+        check_parameter(m, low=1, high= data.num_nodes, param_name='m')
+    else:
+        raise ValueError("m should be int, got %s" % m)
+
+    if isinstance(n, int):
+        check_parameter(n, low=1, high=data.num_nodes, param_name='n')
+    else:
+        raise ValueError("n should be int, got %s" % n)
+
+    check_parameter(m*n, low=1, high=data.num_nodes, param_name='m*n')
+
+    if random_state:
+        np.random.seed(random_state)
+
+    new_edges = []
 
-        sample_indices += sampled_idx.tolist()
+    outlier_idx = np.random.choice(data.num_nodes, size=m * n, replace=False)
 
-        for j, id1 in enumerate(sampled_idx):
-            for k, id2 in enumerate(sampled_idx):
+    # connect all m nodes in each clique
+    for i in range(0, n):
+        for j in range(m * i, m * (i + 1)):
+            for k in range(m * i, m * (i + 1)):
                 if j != k:
+                    node1, node2 = outlier_idx[j], outlier_idx[k]
                     new_edges.append(
-                        torch.tensor([[id1, id2]], dtype=torch.long))
-
-        node_set = node_set.difference(set(sampled_idx))
+                        torch.tensor([[node1, node2]], dtype=torch.long))
 
     y_outlier = torch.zeros(data.x.shape[0], dtype=torch.long)
-    y_outlier[sample_indices] = 1
+    y_outlier[outlier_idx] = 1
 
     data.edge_index = torch.cat([data.edge_index, torch.cat(new_edges).T],
                                 dim=1)
@@ -66,50 +83,68 @@ def gen_structure_outliers(data, m, n):
     return data, y_outlier
 
 
-def gen_attribute_outliers(data, n, k):
+def gen_attribute_outliers(data, n, k, random_state=None):
     """Generating attribute outliers according to paper
     "Deep Anomaly Detection on Attributed Networks"
     <https://epubs.siam.org/doi/abs/10.1137/1.9781611975673.67>.
-
     We randomly select n nodes as the attribute perturbation candidates.
     For each selected node i, we randomly pick another k nodes from the data
     and select the node j whose attributes deviate the most from node i
     among the k nodes by maximizing the Euclidean distance ||xi − xj ||2.
     Afterwards, we then change the attributes xi of node i to xj.
-
     Parameters
     ----------
-    data : torch_geometric.datasets.Data
-        PyG dataset object.
+    data : PyTorch Geometric Data instance (torch_geometric.data.Data)
+        The input data.
     n : int
         Number of nodes converting to outliers.
     k : int
         Number of candidate nodes for each outlier node.
+    random_state : int, optional
+        The seed to control the randomness, Default: ``None``.
+
     Returns
     -------
-    data : torch_geometric.datasets.Data
+    data : PyTorch Geometric Data instance (torch_geometric.data.Data)
         The attribute outlier graph with modified node attributes.
     y_outlier : torch.Tensor
         The outlier label tensor where 1 represents outliers and 0 represents
         regular nodes.
     """
 
-    node_set = set(range(data.x.shape[0]))
-    target_nodes = np.random.choice(list(node_set), size=n, replace=False)
-    candidate_set = node_set.difference(set(target_nodes))
+    if not isinstance(data, Data):
+        raise TypeError("data should be torch_geometric.data.Data")
+
+    if isinstance(n, int):
+        check_parameter(n, low=1, high=data.num_nodes, param_name='n')
+    else:
+        raise ValueError("n should be int, got %s" % n)
+
+    if isinstance(k, int):
+        check_parameter(k, low=1, high=data.num_nodes-n, param_name='k')
+    else:
+        raise ValueError("k should be int, got %s" % k)
+
+    if random_state:
+        np.random.seed(random_state)
+
+    node_set = set(range(data.num_nodes))
+
+    outlier_idx = np.random.choice(list(node_set), size=n, replace=False)
+    candidate_set = node_set.difference(set(outlier_idx))
+    candidate_idx = np.random.choice(list(candidate_set), size=n * k)
 
-    for u in target_nodes:
-        temp_candidates = np.random.choice(list(candidate_set), size=k,
-                                           replace=False)
+    for i, idx in enumerate(outlier_idx):
+        cur_candidates = candidate_idx[k * i: k * (i + 1)]
 
-        euclidean_dist = torch.cdist(data.x[u].unsqueeze(0), data.x[list(
-            temp_candidates)])
+        euclidean_dist = torch.cdist(data.x[idx].unsqueeze(0), data.x[list(
+            cur_candidates)])
         max_dist_idx = torch.argmax(euclidean_dist)
-        max_dist_node = list(temp_candidates)[max_dist_idx]
+        max_dist_node = list(cur_candidates)[max_dist_idx]
 
-        data.x[u] = data.x[max_dist_node]
+        data.x[idx] = data.x[max_dist_node]
 
     y_outlier = torch.zeros(data.x.shape[0], dtype=torch.long)
-    y_outlier[target_nodes] = 1
+    y_outlier[outlier_idx] = 1
 
     return data, y_outlier
diff --git a/pygod/utils/utility.py b/pygod/utils/utility.py
new file mode 100644
index 0000000..e3b7d3c
--- /dev/null
+++ b/pygod/utils/utility.py
@@ -0,0 +1,127 @@
+# -*- coding: utf-8 -*-
+"""A set of utility functions to support outlier detection.
+"""
+# Author: Yue Zhao <zhaoy@cmu.edu>
+# License: BSD 2 clause
+
+from __future__ import division
+from __future__ import print_function
+
+import warnings
+import torch
+import numpy as np
+import numbers
+
+MAX_INT = np.iinfo(np.int32).max
+MIN_INT = -1 * MAX_INT
+
+
+def validate_device(gpu_id):
+    """Validate the input device id (GPU id) is valid on the given
+    machine. If no GPU is presented, return 'cpu'.
+
+    Parameters
+    ----------
+    gpu_id : int
+        GPU id to be used. The function will validate the usability
+        of the GPU. If failed, return device as 'cpu'.
+
+    Returns
+    -------
+    device_id : str
+        Valid device id, e.g., 'cuda:0' or 'cpu'
+    """
+    # if it is cpu
+    if gpu_id == -1:
+        return 'cpu'
+
+    # cast to int for checking
+    gpu_id = int(gpu_id)
+
+    # if gpu is available
+    if torch.cuda.is_available():
+        # check if gpu id is between 0 and the total number of GPUs
+        check_parameter(gpu_id, 0, torch.cuda.device_count(), param_name='gpu id', include_left=True,
+                        include_right=False)
+        device_id = 'cuda:{}'.format(gpu_id)
+    else:
+        if gpu_id != 'cpu':
+            warnings.warn('The cuda is not available. Set to cpu.')
+        device_id = 'cpu'
+
+    return device_id
+
+
+def check_parameter(param, low=MIN_INT, high=MAX_INT, param_name='',
+                    include_left=False, include_right=False):
+    """Check if an input is within the defined range.
+    Parameters
+    ----------
+    param : int, float
+        The input parameter to check.
+    low : int, float
+        The lower bound of the range.
+    high : int, float
+        The higher bound of the range.
+    param_name : str, optional (default='')
+        The name of the parameter.
+    include_left : bool, optional (default=False)
+        Whether includes the lower bound (lower bound <=).
+    include_right : bool, optional (default=False)
+        Whether includes the higher bound (<= higher bound).
+    Returns
+    -------
+    within_range : bool or raise errors
+        Whether the parameter is within the range of (low, high)
+    """
+
+    # param, low and high should all be numerical
+    if not isinstance(param, (numbers.Integral, np.integer, float)):
+        raise TypeError('{param_name} is set to {param} Not numerical'.format(
+            param=param, param_name=param_name))
+
+    if not isinstance(low, (numbers.Integral, np.integer, float)):
+        raise TypeError('low is set to {low}. Not numerical'.format(low=low))
+
+    if not isinstance(high, (numbers.Integral, np.integer, float)):
+        raise TypeError('high is set to {high}. Not numerical'.format(
+            high=high))
+
+    # at least one of the bounds should be specified
+    if low is MIN_INT and high is MAX_INT:
+        raise ValueError('Neither low nor high bounds is undefined')
+
+    # if wrong bound values are used
+    if low > high:
+        raise ValueError(
+            'Lower bound > Higher bound')
+
+    # value check under different bound conditions
+    if (include_left and include_right) and (param < low or param > high):
+        raise ValueError(
+            '{param_name} is set to {param}. '
+            'Not in the range of [{low}, {high}].'.format(
+                param=param, low=low, high=high, param_name=param_name))
+
+    elif (include_left and not include_right) and (
+            param < low or param >= high):
+        raise ValueError(
+            '{param_name} is set to {param}. '
+            'Not in the range of [{low}, {high}).'.format(
+                param=param, low=low, high=high, param_name=param_name))
+
+    elif (not include_left and include_right) and (
+            param <= low or param > high):
+        raise ValueError(
+            '{param_name} is set to {param}. '
+            'Not in the range of ({low}, {high}].'.format(
+                param=param, low=low, high=high, param_name=param_name))
+
+    elif (not include_left and not include_right) and (
+            param <= low or param >= high):
+        raise ValueError(
+            '{param_name} is set to {param}. '
+            'Not in the range of ({low}, {high}).'.format(
+                param=param, low=low, high=high, param_name=param_name))
+    else:
+        return True
diff --git a/pygod/version.py b/pygod/version.py
index b361ab9..493eec6 100644
--- a/pygod/version.py
+++ b/pygod/version.py
@@ -20,4 +20,4 @@
 # Dev branch marker is: 'X.Y.dev' or 'X.Y.devN' where N is an integer.
 # 'X.Y.dev0' is the canonical version of 'X.Y.dev'
 #
-__version__ = '0.1.1'  # pragma: no cover
+__version__ = '0.2.0'  # pragma: no cover
diff --git a/requirements.txt b/requirements.txt
index d0e3fa3..147ae5e 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,4 +1,3 @@
-argparse>=1.4.0
 numpy>=1.19.4
 scikit-learn>=0.22.1
 scipy>=1.5.2
diff --git a/requirements_ci.txt b/requirements_ci.txt
index 0dc0601..1b378a6 100644
--- a/requirements_ci.txt
+++ b/requirements_ci.txt
@@ -1,4 +1,3 @@
-argparse>=1.4.0
 networkx
 numpy>=1.19.4
 scikit-learn>=0.22.1


Test patch:
diff --git a/pygod/test/test_anemone.py b/pygod/test/test_anemone.py
new file mode 100644
index 0000000..ae3b20d
--- /dev/null
+++ b/pygod/test/test_anemone.py
@@ -0,0 +1,109 @@
+# -*- coding: utf-8 -*-
+import os
+import unittest
+# noinspection PyProtectedMember
+from numpy.testing import assert_equal
+from numpy.testing import assert_raises
+
+import torch
+from torch_geometric.seed import seed_everything
+from pygod.models import ANEMONE
+from pygod.utils.metric import eval_roc_auc
+
+seed_everything(42)
+
+
+class TestAnemone(unittest.TestCase):
+    def setUp(self):
+        # use the pre-defined fake graph with injected outliers
+        # for testing purpose
+
+        # the roc should be higher than this; it is model dependent
+        self.roc_floor = 0.5
+
+        test_graph = torch.load(os.path.join('pygod', 'test', 'test_graph.pt'))
+        self.data = test_graph
+
+        self.model = ANEMONE()
+        self.model.fit(self.data)
+
+    def test_parameters(self):
+        assert (hasattr(self.model, 'decision_scores_') and
+                self.model.decision_scores_ is not None)
+        assert (hasattr(self.model, 'labels_') and
+                self.model.labels_ is not None)
+        assert (hasattr(self.model, 'threshold_') and
+                self.model.threshold_ is not None)
+        assert (hasattr(self.model, '_mu') and
+                self.model._mu is not None)
+        assert (hasattr(self.model, '_sigma') and
+                self.model._sigma is not None)
+        assert (hasattr(self.model, 'model') and
+                self.model.model is not None)
+
+    def test_train_scores(self):
+        assert_equal(len(self.model.decision_scores_), len(self.data.y))
+
+    def test_prediction_scores(self):
+        pred_scores = self.model.decision_function(self.data)
+
+        # check score shapes
+        assert_equal(pred_scores.shape[0], self.data.y.shape[0])
+
+        # check performance
+        assert (eval_roc_auc(self.data.y, pred_scores) >= self.roc_floor)
+
+    def test_prediction_labels(self):
+        pred_labels = self.model.predict(self.data)
+        assert_equal(pred_labels.shape[0], self.data.y.shape[0])
+
+    def test_prediction_proba(self):
+        pred_proba = self.model.predict_proba(self.data)
+        assert (pred_proba.min() >= 0)
+        assert (pred_proba.max() <= 1)
+
+    def test_prediction_proba_linear(self):
+        pred_proba = self.model.predict_proba(self.data, method='linear')
+        assert (pred_proba.min() >= 0)
+        assert (pred_proba.max() <= 1)
+
+    def test_prediction_proba_unify(self):
+        pred_proba = self.model.predict_proba(self.data, method='unify')
+        assert (pred_proba.min() >= 0)
+        assert (pred_proba.max() <= 1)
+
+    def test_prediction_proba_parameter(self):
+        with assert_raises(ValueError):
+            self.model.predict_proba(self.data, method='something')
+
+    def test_prediction_labels_confidence(self):
+        pred_labels, confidence = self.model.predict(self.data,
+                                                     return_confidence=True)
+        assert_equal(pred_labels.shape[0], self.data.y.shape[0])
+        assert_equal(confidence.shape[0], self.data.y.shape[0])
+        assert (confidence.min() >= 0)
+        assert (confidence.max() <= 1)
+
+    def test_prediction_proba_linear_confidence(self):
+        pred_proba, confidence = self.model.predict_proba(self.data,
+                                                          method='linear',
+                                                          return_confidence=True)
+        assert (pred_proba.min() >= 0)
+        assert (pred_proba.max() <= 1)
+
+        assert_equal(confidence.shape[0], self.data.y.shape[0])
+        assert (confidence.min() >= 0)
+        assert (confidence.max() <= 1)
+
+    def test_model_clone(self):
+        pass
+        # clone_clf = clone(self.model)
+
+    def tearDown(self):
+        pass
+        # remove the data folder
+        # rmtree(self.path)
+
+
+if __name__ == '__main__':
+    unittest.main()
diff --git a/pygod/test/test_cola.py b/pygod/test/test_cola.py
new file mode 100644
index 0000000..964d540
--- /dev/null
+++ b/pygod/test/test_cola.py
@@ -0,0 +1,109 @@
+# -*- coding: utf-8 -*-
+import os
+import unittest
+# noinspection PyProtectedMember
+from numpy.testing import assert_equal
+from numpy.testing import assert_raises
+
+import torch
+from torch_geometric.seed import seed_everything
+from pygod.models import CoLA
+from pygod.utils.metric import eval_roc_auc
+
+seed_everything(42)
+
+
+class TestCoLA(unittest.TestCase):
+    def setUp(self):
+        # use the pre-defined fake graph with injected outliers
+        # for testing purpose
+
+        # the roc should be higher than this; it is model dependent
+        self.roc_floor = 0.3
+
+        test_graph = torch.load(os.path.join('pygod', 'test', 'test_graph.pt'))
+        self.data = test_graph
+
+        self.model = CoLA()
+        self.model.fit(self.data)
+
+    def test_parameters(self):
+        assert (hasattr(self.model, 'decision_scores_') and
+                self.model.decision_scores_ is not None)
+        assert (hasattr(self.model, 'labels_') and
+                self.model.labels_ is not None)
+        assert (hasattr(self.model, 'threshold_') and
+                self.model.threshold_ is not None)
+        assert (hasattr(self.model, '_mu') and
+                self.model._mu is not None)
+        assert (hasattr(self.model, '_sigma') and
+                self.model._sigma is not None)
+        assert (hasattr(self.model, 'model') and
+                self.model.model is not None)
+
+    def test_train_scores(self):
+        assert_equal(len(self.model.decision_scores_), len(self.data.y))
+
+    def test_prediction_scores(self):
+        pred_scores = self.model.decision_function(self.data)
+
+        # check score shapes
+        assert_equal(pred_scores.shape[0], self.data.y.shape[0])
+
+        # check performance
+        assert (eval_roc_auc(self.data.y, pred_scores) >= self.roc_floor)
+
+    def test_prediction_labels(self):
+        pred_labels = self.model.predict(self.data)
+        assert_equal(pred_labels.shape[0], self.data.y.shape[0])
+
+    def test_prediction_proba(self):
+        pred_proba = self.model.predict_proba(self.data)
+        assert (pred_proba.min() >= 0)
+        assert (pred_proba.max() <= 1)
+
+    def test_prediction_proba_linear(self):
+        pred_proba = self.model.predict_proba(self.data, method='linear')
+        assert (pred_proba.min() >= 0)
+        assert (pred_proba.max() <= 1)
+
+    def test_prediction_proba_unify(self):
+        pred_proba = self.model.predict_proba(self.data, method='unify')
+        assert (pred_proba.min() >= 0)
+        assert (pred_proba.max() <= 1)
+
+    def test_prediction_proba_parameter(self):
+        with assert_raises(ValueError):
+            self.model.predict_proba(self.data, method='something')
+
+    def test_prediction_labels_confidence(self):
+        pred_labels, confidence = self.model.predict(self.data,
+                                                     return_confidence=True)
+        assert_equal(pred_labels.shape[0], self.data.y.shape[0])
+        assert_equal(confidence.shape[0], self.data.y.shape[0])
+        assert (confidence.min() >= 0)
+        assert (confidence.max() <= 1)
+
+    def test_prediction_proba_linear_confidence(self):
+        pred_proba, confidence = self.model.predict_proba(self.data,
+                                                          method='linear',
+                                                          return_confidence=True)
+        assert (pred_proba.min() >= 0)
+        assert (pred_proba.max() <= 1)
+
+        assert_equal(confidence.shape[0], self.data.y.shape[0])
+        assert (confidence.min() >= 0)
+        assert (confidence.max() <= 1)
+
+    def test_model_clone(self):
+        pass
+        # clone_clf = clone(self.model)
+
+    def tearDown(self):
+        pass
+        # remove the data folder
+        # rmtree(self.path)
+
+
+if __name__ == '__main__':
+    unittest.main()
diff --git a/pygod/test/test_conad.py b/pygod/test/test_conad.py
new file mode 100644
index 0000000..bd69e10
--- /dev/null
+++ b/pygod/test/test_conad.py
@@ -0,0 +1,109 @@
+# -*- coding: utf-8 -*-
+import os
+import unittest
+# noinspection PyProtectedMember
+from numpy.testing import assert_equal
+from numpy.testing import assert_raises
+
+import torch
+from torch_geometric.seed import seed_everything
+from pygod.models import CONAD
+from pygod.utils.metric import eval_roc_auc
+
+seed_everything(42)
+
+
+class TestCONAD(unittest.TestCase):
+    def setUp(self):
+        # use the pre-defined fake graph with injected outliers
+        # for testing purpose
+
+        # the roc should be higher than this; it is model dependent
+        self.roc_floor = 0.60
+
+        test_graph = torch.load(os.path.join('pygod', 'test', 'test_graph.pt'))
+        self.data = test_graph
+
+        self.model = CONAD()
+        self.model.fit(self.data)
+
+    def test_parameters(self):
+        assert (hasattr(self.model, 'decision_scores_') and
+                self.model.decision_scores_ is not None)
+        assert (hasattr(self.model, 'labels_') and
+                self.model.labels_ is not None)
+        assert (hasattr(self.model, 'threshold_') and
+                self.model.threshold_ is not None)
+        assert (hasattr(self.model, '_mu') and
+                self.model._mu is not None)
+        assert (hasattr(self.model, '_sigma') and
+                self.model._sigma is not None)
+        assert (hasattr(self.model, 'model') and
+                self.model.model is not None)
+
+    def test_train_scores(self):
+        assert_equal(len(self.model.decision_scores_), len(self.data.y))
+
+    def test_prediction_scores(self):
+        pred_scores = self.model.decision_function(self.data)
+
+        # check score shapes
+        assert_equal(pred_scores.shape[0], self.data.y.shape[0])
+
+        # check performance
+        assert (eval_roc_auc(self.data.y, pred_scores) >= self.roc_floor)
+
+    def test_prediction_labels(self):
+        pred_labels = self.model.predict(self.data)
+        assert_equal(pred_labels.shape[0], self.data.y.shape[0])
+
+    def test_prediction_proba(self):
+        pred_proba = self.model.predict_proba(self.data)
+        assert (pred_proba.min() >= 0)
+        assert (pred_proba.max() <= 1)
+
+    def test_prediction_proba_linear(self):
+        pred_proba = self.model.predict_proba(self.data, method='linear')
+        assert (pred_proba.min() >= 0)
+        assert (pred_proba.max() <= 1)
+
+    def test_prediction_proba_unify(self):
+        pred_proba = self.model.predict_proba(self.data, method='unify')
+        assert (pred_proba.min() >= 0)
+        assert (pred_proba.max() <= 1)
+
+    def test_prediction_proba_parameter(self):
+        with assert_raises(ValueError):
+            self.model.predict_proba(self.data, method='something')
+
+    def test_prediction_labels_confidence(self):
+        pred_labels, confidence = self.model.predict(self.data,
+                                                     return_confidence=True)
+        assert_equal(pred_labels.shape[0], self.data.y.shape[0])
+        assert_equal(confidence.shape[0], self.data.y.shape[0])
+        assert (confidence.min() >= 0)
+        assert (confidence.max() <= 1)
+
+    def test_prediction_proba_linear_confidence(self):
+        pred_proba, confidence = self.model.predict_proba(self.data,
+                                                          method='linear',
+                                                          return_confidence=True)
+        assert (pred_proba.min() >= 0)
+        assert (pred_proba.max() <= 1)
+
+        assert_equal(confidence.shape[0], self.data.y.shape[0])
+        assert (confidence.min() >= 0)
+        assert (confidence.max() <= 1)
+
+    def test_model_clone(self):
+        pass
+        # clone_clf = clone(self.model)
+
+    def tearDown(self):
+        pass
+        # remove the data folder
+        # rmtree(self.path)
+
+
+if __name__ == '__main__':
+    unittest.main()
\ No newline at end of file
diff --git a/pygod/test/test_outlier_generator.py b/pygod/test/test_outlier_generator.py
new file mode 100644
index 0000000..213a337
--- /dev/null
+++ b/pygod/test/test_outlier_generator.py
@@ -0,0 +1,124 @@
+# -*- coding: utf-8 -*-
+import unittest
+# noinspection PyProtectedMember
+from numpy.testing import assert_allclose
+from numpy.testing import assert_equal
+from numpy.testing import assert_raises
+
+import os
+import copy as cp
+
+import torch
+from torch_geometric.seed import seed_everything
+from pygod.utils.outlier_generator import gen_attribute_outliers
+from pygod.utils.outlier_generator import gen_structure_outliers
+
+seed_everything(42)
+
+
+class TestData(unittest.TestCase):
+    def setUp(self):
+        self.m_structure = 5
+        self.n_structure = 5
+        self.n_attribute = 25
+        self.k_attribute = 25
+        self.random_state = 42
+
+        test_graph = torch.load(os.path.join('pygod', 'test', 'test_graph.pt'))
+        self.data = test_graph
+
+    def test_structure_outliers(self):
+        test_data = cp.deepcopy(self.data)
+        test_data, y_outlier = gen_structure_outliers(data=test_data,
+                                                      m=self.m_structure,
+                                                      n=self.n_structure)
+
+        assert_equal(self.data.x.numpy(), test_data.x.numpy())
+        assert (self.data.num_nodes == y_outlier.shape[0])
+        assert (test_data.num_edges ==
+                self.m_structure * (self.m_structure - 1) * self.n_structure
+                + self.data.num_edges)
+        assert_equal(self.data.edge_index.numpy(),
+                     test_data.edge_index[:, :self.data.edge_index.shape[1]])
+        assert (self.m_structure * self.n_structure == y_outlier.sum().item())
+
+    def test_attribute_outliers(self):
+        test_data = cp.deepcopy(self.data)
+        test_data, y_outlier = gen_attribute_outliers(data=test_data,
+                                                      n=self.n_attribute,
+                                                      k=self.k_attribute)
+
+        assert_equal(self.data.edge_index.numpy(),
+                     test_data.edge_index.numpy())
+        assert_allclose(self.data.x.numpy().shape, test_data.x.numpy().shape)
+        assert (self.data.num_nodes == y_outlier.shape[0])
+        assert (self.n_attribute == y_outlier.sum().item())
+
+    def test_structure_outliers2(self):
+        test_data, y_outlier = \
+            gen_structure_outliers(data=cp.deepcopy(self.data),
+                                   m=self.m_structure,
+                                   n=self.n_structure,
+                                   random_state=self.random_state)
+
+        test_data2, y_outlier2 = \
+            gen_structure_outliers(data=cp.deepcopy(self.data),
+                                   m=self.m_structure,
+                                   n=self.n_structure,
+                                   random_state=self.random_state)
+
+        assert_equal(test_data.x.numpy(), test_data2.x.numpy())
+        assert_equal(test_data.edge_index.numpy(),
+                     test_data2.edge_index.numpy())
+        assert_equal(y_outlier.numpy(), y_outlier2.numpy())
+
+    def test_attribute_outliers2(self):
+        test_data, y_outlier = \
+            gen_attribute_outliers(data=cp.deepcopy(self.data),
+                                   n=self.n_attribute,
+                                   k=self.k_attribute,
+                                   random_state=self.random_state)
+
+        test_data2, y_outlier2 = \
+            gen_attribute_outliers(data=cp.deepcopy(self.data),
+                                   n=self.n_attribute,
+                                   k=self.k_attribute,
+                                   random_state=self.random_state)
+
+        assert_equal(test_data.x.numpy(), test_data2.x.numpy())
+        assert_equal(test_data.edge_index.numpy(),
+                     test_data2.edge_index.numpy())
+        assert_equal(y_outlier.numpy(), y_outlier2.numpy())
+
+    def test_structure_outliers3(self):
+        test_data = cp.deepcopy(self.data)
+
+        with assert_raises(ValueError):
+            gen_structure_outliers(data=test_data,
+                                   m='not int',
+                                   n=self.n_structure)
+
+        with assert_raises(ValueError):
+            gen_structure_outliers(data=test_data,
+                                   m=self.m_structure,
+                                   n='not int')
+
+    def test_attribute_outliers3(self):
+        test_data = cp.deepcopy(self.data)
+
+        with assert_raises(ValueError):
+            gen_attribute_outliers(data=test_data,
+                                   n='not int',
+                                   k=self.k_attribute)
+
+        with assert_raises(ValueError):
+            gen_attribute_outliers(data=test_data,
+                                   n=self.n_attribute,
+                                   k='not int')
+
+    def tearDown(self):
+        pass
+
+
+if __name__ == '__main__':
+    unittest.main()
diff --git a/pygod/test/test_utility.py b/pygod/test/test_utility.py
new file mode 100644
index 0000000..f28849f
--- /dev/null
+++ b/pygod/test/test_utility.py
@@ -0,0 +1,108 @@
+# -*- coding: utf-8 -*-
+
+from __future__ import division
+from __future__ import print_function
+
+import os
+import sys
+
+import unittest
+
+# noinspection PyProtectedMember
+import torch
+
+from numpy.testing import assert_equal
+from numpy.testing import assert_raises
+
+# temporary solution for relative imports in case pyod is not installed
+# if pyod is installed, no need to use the following line
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
+
+from pygod.utils.utility import check_parameter
+from pygod.utils.utility import validate_device
+
+
+class TestValidateDevice(unittest.TestCase):
+    def setUp(self):
+        pass
+
+    def test_check_parameter_range(self):
+        assert (validate_device(-1) == 'cpu')
+
+        if torch.cuda.is_available():
+            with assert_raises(ValueError):
+                validate_device(9999)
+
+        if torch.cuda.is_available():
+            assert (validate_device(0) == 'cuda:0')
+            assert (validate_device(torch.cuda.device_count() - 1) == 'cuda:{}'.format(torch.cuda.device_count() - 1))
+
+
+class TestParameters(unittest.TestCase):
+    def setUp(self):
+        pass
+
+    def test_check_parameter_range(self):
+        # verify parameter type correction
+        with assert_raises(TypeError):
+            check_parameter('f', 0, 100)
+
+        with assert_raises(TypeError):
+            check_parameter(1, 'f', 100)
+
+        with assert_raises(TypeError):
+            check_parameter(1, 0, 'f')
+
+        # with assert_raises(TypeError):
+        #     check_parameter(argmaxn(value_list=[1, 2, 3], n=1), 0, 100)
+
+        # if low and high are both unset
+        with assert_raises(ValueError):
+            check_parameter(50)
+
+        # if low <= high
+        with assert_raises(ValueError):
+            check_parameter(50, 100, 99)
+
+        with assert_raises(ValueError):
+            check_parameter(50, 100, 100)
+
+        # check one side
+        with assert_raises(ValueError):
+            check_parameter(50, low=100)
+        with assert_raises(ValueError):
+            check_parameter(50, high=0)
+
+        assert_equal(True, check_parameter(50, low=10))
+        assert_equal(True, check_parameter(50, high=100))
+
+        # if check fails
+        with assert_raises(ValueError):
+            check_parameter(-1, 0, 100)
+
+        with assert_raises(ValueError):
+            check_parameter(101, 0, 100)
+
+        with assert_raises(ValueError):
+            check_parameter(0.5, 0.2, 0.3)
+
+        # if check passes
+        assert_equal(True, check_parameter(50, 0, 100))
+
+        assert_equal(True, check_parameter(0.5, 0.1, 0.8))
+
+        # if includes left or right bounds
+        with assert_raises(ValueError):
+            check_parameter(100, 0, 100, include_left=False,
+                            include_right=False)
+        assert_equal(True, check_parameter(0, 0, 100, include_left=True,
+                                           include_right=False))
+        assert_equal(True, check_parameter(0, 0, 100, include_left=True,
+                                           include_right=True))
+        assert_equal(True, check_parameter(100, 0, 100, include_left=False,
+                                           include_right=True))
+        assert_equal(True, check_parameter(100, 0, 100, include_left=True,
+                                           include_right=True))
+
+    def tearDown(self):
+        pass
diff --git a/pygod/test/test_utils_metric.py b/pygod/test/test_utils_metric.py
index 83f67ff..7089e5a 100644
--- a/pygod/test/test_utils_metric.py
+++ b/pygod/test/test_utils_metric.py
@@ -1,29 +1,25 @@
 # -*- coding: utf-8 -*-
-import os
-import sys
-
 import unittest
 
 # noinspection PyProtectedMember
 from numpy.testing import assert_allclose
-from numpy.testing import assert_equal
-from numpy.testing import assert_raises
 from sklearn.metrics import roc_auc_score
 from sklearn.metrics import precision_score
 from sklearn.metrics import recall_score
+from sklearn.metrics import average_precision_score
 
 from pygod.utils.metric import eval_roc_auc
 from pygod.utils.metric import eval_recall_at_k
 from pygod.utils.metric import eval_precision_at_k
-
-import numpy as np
+from pygod.utils.metric import eval_average_precision
 
 
 class TestMetric(unittest.TestCase):
 
     def setUp(self):
         self.y = [0, 0, 1, 1, 1, 0, 0, 0, 1, 0]
-        self.decision_scores_ = [0.1, 0.2, 0.2, 0.8, 0.2, 0.5, 0.7, 0.9, 1, 0.3]
+        self.decision_scores_ = [0.1, 0.2, 0.2, 0.8, 0.2, 0.5, 0.7, 0.9, 1,
+                                 0.3]
         self.manual_labels = [0, 0, 0, 1, 0, 0, 1, 1, 1, 0]
 
     def test_eval_roc_auc(self):
@@ -36,4 +32,10 @@ def test_eval_recall_at_k(self):
 
     def test_eval_precision_at_k(self):
         assert_allclose(precision_score(self.y, self.manual_labels),
-                        eval_precision_at_k(self.y, self.decision_scores_, k=4))
+                        eval_precision_at_k(self.y,
+                                            self.decision_scores_,
+                                            k=4))
+
+    def test_eval_average_precision(self):
+        assert_allclose(average_precision_score(self.y, self.decision_scores_),
+                        eval_average_precision(self.y, self.decision_scores_))