Improve code

Signed-off-by: YohannL <yohann.lemoine.enst@gmail.com>