From 5bfe0982e93cbc3ee3e7e1afbbd66e3d445f5bdd Mon Sep 17 00:00:00 2001
From: Yaohui Liu <AsakusaRinne@gmail.com>
Date: Fri, 16 Jun 2023 16:15:27 +0800
Subject: [PATCH] feat: add exception catch to code generator.

---
 tools/Tensorflow.CodeGen/FunctionGenerator.cs | 4 ++++
 1 file changed, 4 insertions(+)

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