[GitHub] [tvm] jroesch commented on a change in pull request #7026: [BYOC][TRT] Support batch norm for all ranks <=5, and all axes

[GitBox <gi...@apache.org>]

jroesch commented on a change in pull request #7026:
URL: https://github.com/apache/tvm/pull/7026#discussion_r537154472



##########
File path: src/runtime/contrib/tensorrt/tensorrt_ops.cc
##########
@@ -386,8 +386,35 @@ class BatchNormOpConverter : public TensorRTOpConverter {
     const int axis = std::stoi(params->node.GetAttr<std::vector<std::string>>("axis")[0]);
     const bool scale = std::stoi(params->node.GetAttr<std::vector<std::string>>("scale")[0]);
     const bool center = std::stoi(params->node.GetAttr<std::vector<std::string>>("center")[0]);
-    ICHECK(axis == 1 || axis == 3);
-    const bool need_transpose = axis == 3;
+    auto input_dims = TrtDimsToVector(input->getDimensions());
+    const size_t min_rank = TRT_HAS_IMPLICIT_BATCH(params) ? 3 : 4;
+    const size_t max_rank = TRT_HAS_IMPLICIT_BATCH(params) ? 4 : 5;
+    ICHECK_LE(input_dims.size(), max_rank);

Review comment:
       Could you convert these checks to use `Diagnostic` instead of generating an assertion, we should strive to replace most of these with end-user readable errors. 




----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org