Merge pull request #278 from vbenni/mobilenetv2_main

(nnapi) Added support for zero input dimension for Gemm op
webmachinelearning · Jun 27, 2022 · 1277e11 · 1277e11
2 parents d4318ce + 12dcb4e
commit 1277e11
Showing 1 changed file with 7 additions and 3 deletions.
diff --git a/src/webnn/native/nnapi/GraphNnapi.h b/src/webnn/native/nnapi/GraphNnapi.h
@@ -144,10 +144,14 @@ namespace webnn::native::nnapi {
                                                 const void* buffer = nullptr) {
             std::shared_ptr<NodeInfo> node = std::make_shared<NodeInfo>();
             node->type = desc->type;
-            for (size_t i = 0; i < desc->dimensionsCount; i++) {
-                node->dimensions.push_back(static_cast<uint32_t>(desc->dimensions[i]));
+            if (desc->dimensionsCount == 0) {
+                node->dimensions.push_back(static_cast<uint32_t>(1));
+            } else {
+                for (size_t i = 0; i < desc->dimensionsCount; i++) {
+                    node->dimensions.push_back(static_cast<uint32_t>(desc->dimensions[i]));
+                }
             }
-
+            
             MaybeError error;
             if (buffer) {
                 error = mNnapiMgr->CreateOperandAndSetMemory(name, node, buffer);