code: def __init__(
      self, model, device=None, avg_fn=None, rm_optimizer=False
  ):
      super(AveragedModel, self).__init__()
      self.start_step = -1
      self.end_step = -1
      model = self.filter(model)
      self.module = deepcopy(model)
      self.module.zero_grad(set_to_none=True)
      if rm_optimizer:
          for k, v in vars(self.module).items():
              if isinstance(v, torch.optim.Optimizer):
                  setattr(self.module, k, None)
                  #  print(f"{k} -> {getattr(self.module, k)}")
      if device is not None:
          self.module = self.module.to(device)
      self.register_buffer(
          "n_averaged", torch.tensor(0, dtype=torch.long, device=device)
      )
      if avg_fn is None:
  
          def avg_fn(
              averaged_model_parameter, model_parameter, num_averaged
          ):
              return averaged_model_parameter + (
                  model_parameter - averaged_model_parameter
              ) / (num_averaged + 1)
  
      self.avg_fn = avg_fn
  
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 2 to create a new AveragedModel object from inherited base classes: Module with example: obj = AveragedModel(model, device, avg_fn, rm_optimizer)