code: def update_bn(iterator, model, n_steps, device="cuda"):
      momenta = {}
      for module in model.modules():
          if isinstance(module, torch.nn.modules.batchnorm._BatchNorm):
              module.running_mean = torch.zeros_like(
                  module.running_mean
              )
              module.running_var = torch.ones_like(module.running_var)
              momenta[module] = module.momentum
  
      if not momenta:
          return
  
      was_training = model.training
      model.train()
      for module in momenta.keys():
          module.momentum = None
          module.num_batches_tracked *= 0
  
      #  for input in loader:
      for i in range(n_steps):
          # batches_dictlist: [{env0_data_key: tensor, env0_...}, env1_..., ...]
          batches_dictlist = next(iterator)
          x = torch.cat([dic["x"] for dic in batches_dictlist])
          x = x.to(device)
  
          model(x)
  
      for bn_module in momenta.keys():
          bn_module.momentum = momenta[bn_module]
      model.train(was_training)
  
imports: import copy
  import warnings
  import math
  from copy import deepcopy
  import torch
  from torch.nn import Module
  from torch.optim.lr_scheduler import _LRScheduler
  from domainbed.ur_networks import URResNet
label_desc: Use image node_id 1 for calling a global function with example usage: update_bn(iterator, model, n_steps, device) without return types