diff --git a/pertbio/pertbio/__init__.py b/pertbio/pertbio/__init__.py
index fa2535f..8198836 100644
--- a/pertbio/pertbio/__init__.py
+++ b/pertbio/pertbio/__init__.py
@@ -1,3 +1,6 @@
+"""
+Import all necessary modules
+"""
 from pertbio.config import Config
 from pertbio.model import *
 from pertbio.kernel import *
diff --git a/pertbio/pertbio/model.py b/pertbio/pertbio/model.py
index 485d04e..964cd6e 100644
--- a/pertbio/pertbio/model.py
+++ b/pertbio/pertbio/model.py
@@ -1,6 +1,6 @@
 """
 This module defines the structures for different models, including
-CellBox, linear regression, and co-expression
+CellBox, linear regression, Neural network, and co-expression
 """
 
 import numpy as np
@@ -183,6 +183,7 @@ def forward(self, mu):
 
 
 class NN(LinReg):
+    """Neural network model"""
     def get_variables(self):
         with tf.compat.v1.variable_scope("initialization", reuse=True):
             self.params.update({