diff --git a/src/finn/core/utils.py b/src/finn/core/utils.py index 12c34d8a56b95be363e7929cd4812c8f7981e700..05e3adf1f0c2643c21a1185c7a3c681ba28fd845 100644 --- a/src/finn/core/utils.py +++ b/src/finn/core/utils.py @@ -185,8 +185,10 @@ def gen_finn_dt_tensor(finn_dt, tensor_shape): tensor_values = np.random.randint(2, size=tensor_shape) elif finn_dt == DataType.TERNARY: tensor_values = np.random.randint(-1, high=1, size=tensor_shape) - elif finn_dt == DataType.INT2: - tensor_values = np.random.randint(-2, high=1, size=tensor_shape) + elif "INT" in finn_dt.name: + tensor_values = np.random.randint( + finn_dt.min(), high=finn_dt.max(), size=tensor_shape + ) else: raise ValueError( "Datatype {} is not supported, no tensor could be generated".format(finn_dt)