torch.func.replace_all_batch_norm_modules_#

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

In place updatesroot by setting therunning_mean andrunning_var to be None andsetting track_running_stats to be False for any nn.BatchNorm module inroot

Return type:

Module