Add code to get where to slice the output of the mlp

Signed-off-by: Axel Davy <axel.davy@ens.fr>