1 files changed, 2 insertions, 1 deletions

diff --git a/tosa.xml b/tosa.xml
index 7fb641d..6576db1 100644
--- a/tosa.xml
+++ b/tosa.xml
@@ -17,7 +17,8 @@
         <name>ARGMAX</name>
         <arguments>
           <argument category="input" name="input" type="in_t*" shape="shape1">
-            <description>Input tensor with rank from 1 to 4</description>
+            <description>Input tensor</description>
+            <levellimit value="rank(shape1)" limit="MAX_RANK"/>
           </argument>
           <argument category="attribute" name="axis" type="int32_t" shape="-">
             <description>Axis in range from 0 to rank(shape1)-1</description>