class_docstr: Exponential Moving Average over a nn.ModuleDict.
  
  You can switch to the EMA weights temporarily.
code: def __init__(
      self,
      module_dict: nn.ModuleDict,
      decay: float = 0.999,
      unbias: bool = True,
      device: tp.Union[torch.device, str] = "cpu",
  ):
      self.decay = decay
      self.module_dict = module_dict
      self.state: dict = defaultdict(dict)
      self.count = 0
      self.device = device
      self.unbias = unbias
      self._init()
  
file_path: audiocraft/audiocraft/optim/ema.py
imports: from collections import defaultdict
  import typing
  import torch
  import torch.nn
label_desc: Use image node_id 1 to create a new ModuleDictEMA object  with example: obj = ModuleDictEMA(module_dict, decay, unbias, device)