model_info: add some tflite types

tum-ei-eda · Jun 24, 2022 · 7030063 · 7030063
1 parent 59977ff
commit 7030063
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/mlonmcu/flow/tvm/backend/model_info.py b/mlonmcu/flow/tvm/backend/model_info.py
@@ -28,12 +28,13 @@ def __init__(self, name, shape, dtype, fix_names=False):
             self.name = self.name.replace("/", "_").replace(";", "_")
         assert isinstance(shape, (tuple, list))
         self.shape = shape
-        assert dtype in ["float32", "uint8", "int8"]
         size_lookup = {
             "float32": 4,
             "uint8": 1,
             "int8": 1,
+            "int32": 4,
         }
+        assert dtype in size_lookup
         self.dtype = dtype
         self.type_size = size_lookup[self.dtype]
 
@@ -54,6 +55,8 @@ def __init__(self, t, fix_names=False):
             TType.FLOAT32: "float32",
             TType.UINT8: "uint8",
             TType.INT8: "int8",
+            TType.INT32: "int32",
+            TType.BOOL: "int8",
         }
         dtype = type_lookup[t.Type()]
         super().__init__(name, shape, dtype, fix_names=fix_names)