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