BatchNorm don't work as expected

[wangjiawen2013]

Hi,
Here is the output of pytorch BatchNorm1d, it worked well:

and here is the rust burn BatchNorm, the output is weired:

Then I changed to AutodiffBackend, the results is correct ! Why don't limit the BatchNorm to only AutodiffBackend ? Otherwise we are prone to make mistakes.

[laggui]

The result is not wrong 🙂

The difference lies in the training vs inference computation for a batchnorm module. If you use m.eval() instead for the pytorch module you should get equivalent results.

With burn you have to be explicit when using autodiff. In pytorch it's kind of the opposite, it will track the gradients and keep the autodiff graph by default unless you use with torch.no_grad() context.

This is also explained in the autodiff section.