code: def split_tensor_along_last_dim(
      tensor, partitions, contiguous_split_chunks=False
  ):
      """Split a tensor along its last dimension. Adapted from Megatron-LM.
  
      Arguments:
          tensor: input tensor.
          partitions: list of partition sizes to supply to torch.split
          contiguous_split_chunks: If True, make each chunk contiguous
                                   in memory.
      """
      # Get the size and dimension.
      last_dim = tensor.dim() - 1
  
      # Split.
      tensor_list = torch.split(tensor, partitions, dim=last_dim)
      # Note: torch.split does not create contiguous tensors by default.
      if contiguous_split_chunks:
          return tuple(chunk.contiguous() for chunk in tensor_list)
  
      return tensor_list
  
file_path: DeepSpeed/deepspeed/runtime/zero/tiling.py
imports: import torch
  import deepspeed
  from deepspeed.runtime.utils import partition_uniform
label_desc: Use image node_id 1 for calling a global function with example usage: split_tensor_along_last_dim(tensor, partitions, contiguous_split_chunks) and returns: tensor_list, tuple