pytorch-maml | An Implementation of Model-Agnostic Meta-Learning | Machine Learning library

TLDR: Use mdl.train() since that uses batch statistics (but inference will not be deterministic anymore). You probably won't want to use mdl.eval() in meta-learning.

BN intended behaviour:

• Importantly, during inference (eval/testing) running_mean, running_std is used - that was calculated from training(because they want a deterministic output and to use estimates of the population statistics).
• During training the batch statistics is used but a population statistic is estimated with running averages. I assume the reason batch_stats is used during training is to introduce noise that regularizes training (noise robustness)
• in meta-learning I think using batch statistics is the best during testing (and not calculate the running means) since we are supposed to be seeing new /tasksdistribution anyway. Price we pay is loss of determinism. Could be interesting just out of curiosity what the accuracy is using population stats estimated from meta-trian.

This is likely why I don't see divergence in my testing with the mdl.train().

So just make sure you use mdl.train() (since that uses batch statistics https://pytorch.org/docs/stable/generated/torch.nn.BatchNorm2d.html#torch.nn.BatchNorm2d) but that either the new running stats that cheat aren't saved or used later.