update scripts/utils/tvm/infer_type.py

tum-ei-eda · Aug 19, 2024 · 8df0d38 · 8df0d38
1 parent 0740f3a
commit 8df0d38
Showing 1 changed file with 7 additions and 2 deletions.
diff --git a/scripts/utils/tvm/infer_type.py b/scripts/utils/tvm/infer_type.py
@@ -7,7 +7,7 @@
 
 # from tvm import parser
 
-assert len(sys.argv) == 2, "Invalid number of arguments"
+assert len(sys.argv) in [2, 3], "Invalid number of arguments"
 
 with open(sys.argv[1]) as f:
     text = f.read()
@@ -17,4 +17,9 @@
 with tvm.transform.PassContext():
     ir_mod = relay.transform.InferType()(ir_mod)
 
-print(ir_mod.astext())
+text = ir_mod.astext()
+if len(sys.argv) == 3:
+    with open(sys.argv[2], "w") as f:
+        f.write(text)
+else:
+    print(text)