diff --git a/config.pbtxt b/config.pbtxt
index b4baace..6c9ac07 100644
--- a/config.pbtxt
+++ b/config.pbtxt
@@ -26,7 +26,7 @@ input [
   {
     name: "max_new_tokens"
     data_type: TYPE_INT32
-    dims: [ 1 ]
+    dims: [ apple ]
 
     optional: true