diff --git a/benchmark.py b/benchmark.py
index 1bb9df8..5f02dfc 100644
--- a/benchmark.py
+++ b/benchmark.py
@@ -86,6 +86,7 @@ def benchmark(model, strategy, timelimit, result_dir, subsolver="scip"):
         # "modprodnet",
         # "stranded_gas",
         # "syngas",
+        # "water_network",
     ]
     strategy_list = [
         "gdp.bigm",
diff --git a/gdplib/__init__.py b/gdplib/__init__.py
index a3cb501..dca67c2 100644
--- a/gdplib/__init__.py
+++ b/gdplib/__init__.py
@@ -12,3 +12,4 @@
 import gdplib.disease_model
 import gdplib.med_term_purchasing
 import gdplib.syngas
+import gdplib.water_network
diff --git a/gdplib/water_network/__init__ .py b/gdplib/water_network/__init__ .py
deleted file mode 100644
index f77f67b..0000000
--- a/gdplib/water_network/__init__ .py	
+++ /dev/null
@@ -1,18 +0,0 @@
-# from .wn_pwl import build_model_Piecewise as _pwl
-# from .wn_q import build_model_Quadratic as _q
-
-# from .wn_minlp import build_model_MINLP as _original
-from gdplib.water_network.wnd import build_model
-
-# def build_model(case="quadratic"):
-#     if case == "quadratic":
-#         return _q()
-#     elif case == "piecewise":
-#         return _pwl()
-#     else:
-#         raise ValueError(f"Invalid case: {case}")
-
-
-__all__ = ['build_model']
-# if __name__ == "__main__":
-#     build_model
diff --git a/gdplib/water_network/__init__.py b/gdplib/water_network/__init__.py
new file mode 100644
index 0000000..d3ac558
--- /dev/null
+++ b/gdplib/water_network/__init__.py
@@ -0,0 +1,3 @@
+from .wnd import build_model
+
+__all__ = ['build_model']
diff --git a/generate_model_size_report.py b/generate_model_size_report.py
index 96ddac0..5995147 100644
--- a/generate_model_size_report.py
+++ b/generate_model_size_report.py
@@ -20,6 +20,7 @@
         # "modprodnet",
         # "stranded_gas",
         # "syngas",
+        # "water_network"
     ]
     current_time = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
     timelimit = 600