text: import torch
  import numpy as np
  
  def size_splits(tensor, split_sizes, dim=0):
      """Splits the tensor according to chunks of split_sizes.
  
      Arguments:
          tensor (Tensor): tensor to split.
          split_sizes (list(int)): sizes of chunks
          dim (int): dimension along which to split the tensor.
      """
      if dim < 0:
          dim += tensor.dim()
  
      dim_size = tensor.size(dim)
      if dim_size != torch.sum(torch.Tensor(split_sizes)):
          raise KeyError("Sum of split sizes exceeds tensor dim")
  
      splits = torch.cumsum(torch.Tensor([0] + split_sizes), dim=0)[:-1]
  
      return tuple(tensor.narrow(int(dim), int(start), int(length))
                   for start, length in zip(splits, split_sizes))
  