Update graph_rewriter.py (#13)

inisis · Jun 20, 2024 · 860ca7f · 860ca7f
1 parent 83708b8
commit 860ca7f
Showing 1 changed file with 2 additions and 5 deletions.
diff --git a/onnxslim/core/graph_rewriter.py b/onnxslim/core/graph_rewriter.py
@@ -26,18 +26,15 @@ 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
 
 
 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