code: def split_experience_batch(
      experience: Experience,
  ) -> List[BufferItem]:
      batch_size = experience.sequences.size(0)
      batch_kwargs = [{} for _ in range(batch_size)]
      keys = (
          "sequences",
          "action_log_probs",
          "values",
          "reward",
          "advantages",
          "attention_mask",
          "action_mask",
      )
      for key in keys:
          value = getattr(experience, key)
          if isinstance(value, torch.Tensor):
              vals = torch.unbind(value)
          else:
              # None
              vals = [value for _ in range(batch_size)]
          assert batch_size == len(vals)
          for i, v in enumerate(vals):
              batch_kwargs[i][key] = v
      items = [BufferItem(**kwargs) for kwargs in batch_kwargs]
      return items
  
file_path: ColossalAI/applications/Chat/coati/experience_buffer/utils.py
imports: from dataclasses import dataclass
  from typing import List, Optional
  import torch
  import torch.nn.functional
  from coati.experience_maker.base import Experience
label_desc: Use image node_id 1 for calling a global function with example usage: split_experience_batch(experience) and returns: items