rf.BatchNorm keeps updating statistics when used in eval mode in training

[mnghiap]

Currently rf.BatchNorm decides whether to update the running statistics based on the rf.get_run_ctx().train_flag as in this line.

update_running_stats = self.running_mean is not None and rf.get_run_ctx().train_flag

However, there are cases where it should be used in eval mode during training (e.g. the teacher model in knowledge distillation), and in those cases it should not update the running statistics.

I suggest that we add something to allow rf.BatchNorm (and maybe other modules with similar issues, but I don't know yet) to run in eval mode even in training.

[albertz]

I was planning to add some context scope where we can overwrite the train_flag. So usage would be sth like:

with rf.get_run_ctx().set_train_flag_scope(False):
  ...

That should solve this, right?