code: def average_state_dicts(state_dicts: List[Dict[str, torch.Tensor]]):
      new_sd = {}
      for k in state_dicts[0].keys():
          tensors = [sd[k] for sd in state_dicts]
          new_t = sum(tensors) / len(tensors)
          assert isinstance(new_t, torch.Tensor)
          new_sd[k] = new_t
      return new_sd
  
file_path: transformers/examples/research_projects/seq2seq-distillation/convert_pl_checkpoint_to_hf.py
imports: import os
  from pathlib import Path
  from typing import Dict, List
  import fire
  import torch
  from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
  from transformers.utils.logging import get_logger
label_desc: Use image node_id 3 for calling a global function with example usage: average_state_dicts(state_dicts) and returns: new_sd