diff --git a/cmsml/tensorflow/__init__.py b/cmsml/tensorflow/__init__.py
index 9042869..a6fc657 100644
--- a/cmsml/tensorflow/__init__.py
+++ b/cmsml/tensorflow/__init__.py
@@ -5,15 +5,19 @@
 Classes, functions and tools for efficiently working with TensorFlow.
 """
 
-__all__ = ["import_tf", "save_frozen_graph", "load_frozen_graph", "write_graph_summary",
-"load_model", "load_graph_def","OpsData", "get_graph_ops"]
+__all__ = [
+    "import_tf", "save_frozen_graph", "save_graph", "load_frozen_graph", "load_graph",
+    "write_graph_summary", "load_model", "load_graph_def",
+    "OpsData", "get_graph_ops",
+]
 
 
 # provisioning imports
 from cmsml.tensorflow.tools import (
-    import_tf, save_frozen_graph, load_frozen_graph, write_graph_summary, load_model, load_graph_def
+    import_tf, save_frozen_graph, save_graph, load_frozen_graph, load_graph, write_graph_summary,
+    load_model, load_graph_def,
 )
 
 from cmsml.tensorflow.aot import (
-    OpsData, get_graph_ops
+    OpsData, get_graph_ops,
 )