xla

257ddd13 - Fix batch normalization to be able to deal with train() and eval() modes.