Add the BatchNorm train/test flag support

Description

This MR aims to bring the support of the train/test mode flag, for the BatchNorm operator.

Two MRs will follow this one : one for the CPU backend and one for the CUDA backend.