hotfix: tflite scalar tensor parsing error

uTensor · Nov 10, 2020 · 3bacf8b · 3bacf8b
1 parent 7d18e68
commit 3bacf8b
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/utensor_cgen/frontend/tflite.py b/utensor_cgen/frontend/tflite.py
@@ -100,6 +100,9 @@ def _build_tensor_map(self, fb_model, ugraph):
 
       if isinstance(tensor.ShapeAsNumpy(), np.ndarray):
         shape = tensor.ShapeAsNumpy().tolist()
+      elif isinstance(tensor.ShapeAsNumpy(), int):
+        logger.warning(f"{tensor.Name().decode('utf8')} is scalar, convert to tensor as shape [1]")
+        shape = [1]
       else:
         shape = list(fb_model.Buffers(12).DataAsNumpy().view(dtype).shape)