torch.func.replace_all_batch_norm_modules_¶

torch.func.replace_all_batch_norm_modules_(root)[source]¶

通过将 running_mean 和 running_var 设置为 None，并将任何 nn.BatchNorm 模块的 track_running_stats 设置为 False 来实现就地更新 root