address pr comments

microsoft · Nov 22, 2023 · bb23a14 · bb23a14
1 parent 7a1408f
commit bb23a14
Show file tree

Hide file tree

Showing 4 changed files with 6 additions and 7 deletions.
diff --git a/onnxruntime/core/providers/coreml/builders/impl/split_op_builder.cc b/onnxruntime/core/providers/coreml/builders/impl/split_op_builder.cc
@@ -82,7 +82,7 @@ Status SplitOpBuilder::AddToModelBuilderImpl(ModelBuilder& model_builder,
     coreml_splitnd->set_numsplits(num_outputs);
   } else {
     // note: for opset 18+ 'num_outputs' is a required attribute
-    num_outputs = narrow<uint64_t>(helper.Get("num_outputs").value());
+    num_outputs = narrow<uint64_t>(helper.GetInt("num_outputs").value());
     // note: checked in IsOpSupportedImpl that ensures the dim value at splitting axis exists
     auto split_dim_size = data_shape[HandleNegativeAxis(axis, data_shape.size())];
     uint64_t chunk_size = narrow<uint64_t>((split_dim_size + num_outputs - 1) / num_outputs);
@@ -159,7 +159,7 @@ bool SplitOpBuilder::IsOpSupportedImpl(const Node& node, const OpBuilderInputPar
     }
   } else {
     if (node.SinceVersion() >= 18) {
-      const auto num_outputs = helper.Get("num_outputs");
+      const auto num_outputs = helper.GetInt("num_outputs");
       if (!num_outputs.has_value()) {
         LOGS(logger, VERBOSE) << "No 'num_outputs' provided. For split 18+, num_outputs is a required attribute.";
         return false;

diff --git a/onnxruntime/core/providers/shared/utils/utils.cc b/onnxruntime/core/providers/shared/utils/utils.cc
@@ -166,11 +166,10 @@ std::vector<float> NodeAttrHelper::Get(const std::string& key, const std::vector
   return std::vector<float>{source.cbegin(), source.cend()};
 }
 
-std::optional<int32_t> NodeAttrHelper::Get(const std::string& key) const {
+std::optional<int64_t> NodeAttrHelper::GetInt(const std::string& key) const {
   if (!HasAttr(key))
     return std::nullopt;
-
-  return SafeInt<int32_t>(node_attributes_.at(key).f());
+  return node_attributes_.at(key).i();
 }
 
 bool NodeAttrHelper::HasAttr(const std::string& key) const {

diff --git a/onnxruntime/core/providers/shared/utils/utils.h b/onnxruntime/core/providers/shared/utils/utils.h
@@ -58,7 +58,7 @@ class NodeAttrHelper {
   uint32_t Get(const std::string& key, uint32_t def_val) const;
   std::vector<uint32_t> Get(const std::string& key, const std::vector<uint32_t>& def_val) const;
 
-  std::optional<int32_t> Get(const std::string& key) const;
+  std::optional<int64_t> GetInt(const std::string& key) const;
 
   bool HasAttr(const std::string& key) const;
 

diff --git a/onnxruntime/test/providers/cpu/tensor/split_op_test.cc b/onnxruntime/test/providers/cpu/tensor/split_op_test.cc
@@ -127,7 +127,7 @@ void SplitTestAxis0EqualSplit(bool use_opset_13 = false) {
              // TensorRT parser: Assertion failed: axis != BATCH_DIM
              {kTensorrtExecutionProvider},  // is_tensorrt_supported
              false,                         // expect_failure
-             use_opset_13 = true);          // split_as_input
+             true);                         // split_as_input
 #endif
 }