16 bit network support
Change-Id: I65cd027dea115d2f50f302ebe997d6d2525e0d7e
diff --git a/driver_library/src/ethosu.cpp b/driver_library/src/ethosu.cpp
index 6b30827..2c498a8 100644
--- a/driver_library/src/ethosu.cpp
+++ b/driver_library/src/ethosu.cpp
@@ -57,12 +57,25 @@
return size;
}
+size_t getTensorTypeSize(const enum tflite::TensorType type) {
+ switch (type) {
+ case tflite::TensorType::TensorType_UINT8:
+ case tflite::TensorType::TensorType_INT8:
+ return 1;
+ case tflite::TensorType::TensorType_INT16:
+ return 2;
+ default:
+ throw EthosU::Exception("Unsupported tensor type");
+ }
+}
+
vector<size_t> getSubGraphDims(const tflite::SubGraph *subgraph, const flatbuffers::Vector<int32_t> *tensorMap) {
vector<size_t> dims;
for (auto index = tensorMap->begin(); index != tensorMap->end(); ++index) {
auto tensor = subgraph->tensors()->Get(*index);
size_t size = getShapeSize(tensor->shape());
+ size *= getTensorTypeSize(tensor->type());
if (size > 0) {
dims.push_back(size);