diff --git a/verif/generator/tosa_arg_gen.py b/verif/generator/tosa_arg_gen.py
index 5957a33..26dd6f9 100644
--- a/verif/generator/tosa_arg_gen.py
+++ b/verif/generator/tosa_arg_gen.py
@@ -1080,9 +1080,9 @@
         for idx, shape in enumerate(shapeList[:]):
             if idx == 1:
                 if dtypeList[idx] == DType.INT8:
-                    arr = np.int32(rng.integers(low=0, high=8, size=shape))
+                    arr = np.int8(rng.integers(low=0, high=8, size=shape))
                 elif dtypeList[idx] == DType.INT16:
-                    arr = np.int32(rng.integers(low=0, high=16, size=shape))
+                    arr = np.int16(rng.integers(low=0, high=16, size=shape))
                 elif dtypeList[idx] == DType.INT32:
                     arr = np.int32(rng.integers(low=0, high=32, size=shape))
                 elif error_name == ErrorIf.WrongInputType: