diff --git a/tools/Tensorflow.CodeGen/FunctionGenerator.cs b/tools/Tensorflow.CodeGen/FunctionGenerator.cs
index 186e6a27..bb07dddf 100644
--- a/tools/Tensorflow.CodeGen/FunctionGenerator.cs
+++ b/tools/Tensorflow.CodeGen/FunctionGenerator.cs
@@ -83,6 +83,10 @@ namespace Tensorflow.CodeGen
 
                 sb.AppendLine("}"); // try
 
+                sb.Append("catch(NotOkStatusException ex)\n{\n");
+                sb.AppendLine("throw ex;");
+                sb.AppendLine("}"); // catch
+
                 sb.Append("catch(Exception)\n{\n");
                 sb.AppendLine("}"); // catch