code: def transfer_batch_norm_stats(
      src_module: torch.nn.Module, dst_module: torch.nn.Module
  ) -> None:
      """
      Transfer batch norm statistics between two same models
      """
      src_batch_norm_modules = []
      dst_batch_norm_modules = []
  
      # fetch all batch norm modules for both
      for module in src_module.modules():
          if isinstance(module, torch.nn.modules.batchnorm._BatchNorm):
              src_batch_norm_modules.append(module)
  
      for module in dst_module.modules():
          if isinstance(module, torch.nn.modules.batchnorm._BatchNorm):
              dst_batch_norm_modules.append(module)
  
      if len(src_batch_norm_modules) != len(dst_batch_norm_modules):
          raise ValueError(
              "Modules must have same number of batch norm layers"
              f"Src module has {len(src_batch_norm_modules)}"
              f"Dst module has {len(dst_batch_norm_modules)}"
          )
  
      # copy batch norm statistics
      for src_batch_norm_module, dst_batch_norm_module in zip(
          src_batch_norm_modules, dst_batch_norm_modules
      ):
          dst_batch_norm_module.running_mean.detach().copy_(
              src_batch_norm_module.running_mean.to(
                  dst_batch_norm_module.running_mean.device
              )
          )
          dst_batch_norm_module.running_var.detach().copy_(
              src_batch_norm_module.running_var.to(
                  dst_batch_norm_module.running_var.device
              )
          )
          dst_batch_norm_module.num_batches_tracked.detach().copy_(
              src_batch_norm_module.num_batches_tracked.to(
                  dst_batch_norm_module.num_batches_tracked.device
              )
          )
          dst_batch_norm_module.momentum = (
              src_batch_norm_module.momentum
          )
  
imports: from typing import Optional
  import torch
label_desc: Use image node_id 3 for calling a global function with example usage: transfer_batch_norm_stats(src_module, dst_module) without return types