code: def split_output(
      output: Tensor,
      offset_dict: Union[Dict[NodeType, int], Dict[EdgeType, int]],
  ) -> Union[Dict[NodeType, Tensor], Dict[EdgeType, Tensor]]:
      cumsums = list(offset_dict.values()) + [output.size(-2)]
      sizes = [
          cumsums[i + 1] - cumsums[i] for i in range(len(offset_dict))
      ]
      outputs = output.split(sizes, dim=-2)
      return {key: output for key, output in zip(offset_dict, outputs)}
  
imports: import copy
  import warnings
  from typing import Any, Dict, List, Optional, Union
  import torch
  from torch import Tensor
  from torch.nn import Module, Parameter
  from torch_geometric.nn.conv import MessagePassing
  from torch_geometric.nn.dense import Linear
  from torch_geometric.nn.fx import Transformer
  from torch_geometric.typing import EdgeType, Metadata, NodeType, SparseTensor
  from torch_geometric.utils.hetero import get_unused_node_types
label_desc: Use image node_id 7 for calling a global function with example usage: split_output(output, offset_dict) and returns: unknown