torch.func.replace_all_batch_norm_modules_

torch.func.replace_all_batch_norm_modules_(root)

In place updates root by setting the running_mean and running_var to be None and setting track_running_stats to be False for any nn.BatchNorm module in root

Return type

Module