code: def update_ema(
      model_parameters,
      ema_parameters,
      model_buffers,
      ema_buffers,
      momentum=0.995,
  ):
      """Update exponential moving average in parameters and buffers."""
      with torch.no_grad():
          torch._foreach_mul(
              ema_parameters, momentum
          )  # want to prevent a second call here, but doesnt seem possible as of now?
          torch._foreach_add_(
              ema_parameters, model_parameters, alpha=1 - momentum
          )
  
          torch._foreach_mul(ema_buffers, momentum)
          torch._foreach_add_(
              ema_buffers, model_buffers, alpha=1 - momentum
          )
  
imports: import os
  import torch
  import transformers
  from torch.utils.data import DataLoader
  from datasets.distributed import split_dataset_by_node
label_desc: Use image node_id 3 for calling a global function with example usage: update_ema(model_parameters, ema_parameters, model_buffers, ema_buffers, momentum) without return types