BatchGradientVerification should probably turn off batch norm

[indigoviolet]

🐛 Bug

In networks using Batch norm, there is "mixing" across the batch dimension but this probably doesn't indicate an error. Might be worth augmenting BatchGradientVerification and the associated callback to

(1) either always turn off batch norm, OR
(2) if we're about to return False, then check if turning it off will return True

Otherwise it can be super confusing to add this check and have it fail, forgetting about batch norm

To Reproduce

https://colab.research.google.com/drive/1dlY2vCBqMasrihIv_Y4DU-3BeKLd0Gox#scrollTo=fDRd_pRn_jL1

Expected behavior

See above.

cc @awaelchli

[github-actions]

Hi! thanks for your contribution!, great first issue!

[awaelchli]

Hi @indigoviolet
Thanks for trying out my callback hihi <3
That's a very good observation! I was aware of it and also bumped into this myself a few times haha, just never had the time to address it properly. I should have made it clearer that it is a heuristic :) But of course an extension would be nice to support this. Your two proposed solutions are both reasonable imo. I would go with the first one since it's the simplest.

Do you have an interest in sending a PR?

[indigoviolet]

cool, PR attached @awaelchli