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