From 860ca7f48f87adbd3599729c6d4b9ec32304aee4 Mon Sep 17 00:00:00 2001
From: Kayzwer <68285002+Kayzwer@users.noreply.github.com>
Date: Thu, 20 Jun 2024 11:25:20 +0800
Subject: [PATCH] Update graph_rewriter.py (#13)

---
 onnxslim/core/graph_rewriter.py | 7 ++-----
 1 file changed, 2 insertions(+), 5 deletions(-)

diff --git a/onnxslim/core/graph_rewriter.py b/onnxslim/core/graph_rewriter.py
index 8db8311..4f08831 100644
--- a/onnxslim/core/graph_rewriter.py
+++ b/onnxslim/core/graph_rewriter.py
@@ -26,10 +26,7 @@ def get_node_feeds(node):
             feeds.append(input)
         else:
             for feed in input.inputs:
-                if feed.op == "Split":
-                    feeds.append(input)
-                else:
-                    feeds.append(feed)
+                feeds.append(input if feed.op == "Split" else feed)
     return feeds
 
 
@@ -37,7 +34,7 @@ def get_name(name):
     _illegal_char_regex = re.compile("[^0-9a-zA-Z_]+")
     sanitized_name = _illegal_char_regex.sub("_", name)
     if sanitized_name.isdigit():
-        sanitized_name = "_" + sanitized_name
+        sanitized_name = f"_{sanitized_name}"
 
     return sanitized_name