diff options
Diffstat (limited to 'arm_compute/core/TensorShape.h')
-rw-r--r-- | arm_compute/core/TensorShape.h | 7 |
1 files changed, 4 insertions, 3 deletions
diff --git a/arm_compute/core/TensorShape.h b/arm_compute/core/TensorShape.h index fe3921f766..7c5ea8d1b7 100644 --- a/arm_compute/core/TensorShape.h +++ b/arm_compute/core/TensorShape.h @@ -71,11 +71,12 @@ public: * * @param[in] dimension Dimension for which the value is set. * @param[in] value Value to be set for the dimension. - * @param[in] apply_dim_correction Flag to state whether apply dimension correction after setting one dimension. E.g. when permuting NCHW -> NHWC, 1x1x2 would become 2x1x1, but _num_dimensions should be 3 rather than 1. + * @param[in] apply_dim_correction (Optional) Flag to state whether apply dimension correction after setting one dimension. E.g. when permuting NCHW -> NHWC, 1x1x2 would become 2x1x1, but _num_dimensions should be 3 rather than 1. + * @param[in] increase_dim_unit (Optional) Set to true if new unit dimensions increase the number of dimensions of the shape. * * @return *this. */ - TensorShape &set(size_t dimension, size_t value, bool apply_dim_correction = true) + TensorShape &set(size_t dimension, size_t value, bool apply_dim_correction = true, bool increase_dim_unit = true) { // Clear entire shape if one dimension is zero if(value == 0) @@ -90,7 +91,7 @@ public: // Set the specified dimension and increase the number of dimensions if // necessary - Dimensions::set(dimension, value, false); + Dimensions::set(dimension, value, increase_dim_unit); // Correct number dimensions to ignore trailing dimensions of size 1 if(apply_dim_correction) |