torch.func.replace_all_batch_norm_modules_

torch.func.replace_all_batch_norm_modules_(root)

通过将running_mean和running_var设置为None，并将root中任何nn.BatchNorm模块的track_running_stats参数设置为False来进行就地更新。