code: def _extract_non_tensor_state(
      combined_state: Dict[int, Dict[str, List]], param_id: int
  ) -> Dict:
      non_tensor_state = (
          {}
      )  # This state is like the `step` count in Adam, not a tensor so we dont unpad or cat it.
      for k, v in combined_state[param_id].items():
          if torch.is_tensor(v[0]):
              continue
          elif len(set(v)) == 1:
              non_tensor_state[k] = v[0]
          else:
              raise TypeError(
                  f"Dont know how to consolidate optimizer param {k} with values {v}"
              )
      return non_tensor_state
  
file_path: fairscale/fairscale/nn/data_parallel/fsdp_optim_utils.py
imports: import copy
  from itertools import groupby
  from typing import TYPE_CHECKING, Any, Dict, Iterator, List, Tuple, cast
  import torch
  from fairscale.nn.misc import FlattenParamsWrapper
label_desc: Use image node_id 3 for calling a global function with example usage: _extract_non_tensor_state(combined_state, param_id) and returns: non_tensor_state