diff --git a/m1_tf_test.py b/m1_tf_test.py
index 7649545..9344a90 100644
--- a/m1_tf_test.py
+++ b/m1_tf_test.py
@@ -6,11 +6,45 @@
 # Author: Smiril
 # Email: sonar@gmx.com
 #############################################################
+import sys
 import tensorflow as tf
 tf.__version__
 tf.config.list_physical_devices()
+
+def xatoi(Str):
+  
+    sign, base, i = 1, 0, 0
+      
+    # If whitespaces then ignore.
+    while (Str[i] == ' '):
+        i += 1
+      
+    # Sign of number
+    if (Str[i] == '-' or Str[i] == '+'):
+        sign = 1 - 2 * (Str[i] == '-')
+        i += 1
+  
+    # Checking for valid input
+    while (i < len(Str) and 
+          Str[i] >= '0' and Str[i] <= '9'):
+                
+        # Handling overflow test case
+        if (base > (sys.maxsize // 10) or
+           (base == (sys.maxsize // 10) and 
+           (Str[i] - '0') > 7)):
+            if (sign == 1):
+                return sys.maxsize
+            else:
+                return -(sys.maxsize)
+          
+        base = 10 * base + (ord(Str[i]) - ord('0'))
+        i += 1
+      
+    return base * sign
+
+user_input = sys.argv[1]
 logits = [[4.0, 2.0, 1.0], [0.0, 5.0, 1.0]]
-inputs = tf.keras.Input(shape=(784,), name="digits")
+inputs = tf.keras.Input(shape=(xatoi(user_input),), name="digits")
 model = tf.keras.models.load_model('model')
 mnist = tf.keras.datasets.mnist
 (x_train, y_train), (x_test, y_test) = mnist.load_data()