class_docstr: Adapted from
  https://github.com/HazyResearch/flash-attention/blob/main/flash_attn/bert_padding.py
code: def forward(ctx, tensor: torch.Tensor, indices: torch.Tensor):
      ctx.save_for_backward(indices)
      # [b, s, ...]
      assert tensor.ndim >= 3
      ctx.bsz = tensor.shape[0]
      out = rearrange(tensor, "b s ... -> (b s) ...")
      ctx.shape = out.shape
      # [ntokens, ...]
      return out[indices]
  
file_path: ColossalAI/colossalai/kernel/cuda_native/mha/utils.py
imports: from dataclasses import dataclass
  from typing import Iterable, Tuple
  import torch
  import torch.nn.functional
  from einops import rearrange
  from colossalai.utils.device import get_current_device
label_desc: Use image node_id 1 for calling the Unpad obj's underlying member method code with example usage: obj.forward(ctx, tensor, indices) and returns: out