torch.func.replace_all_batch_norm_modules_#

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

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