code: def materialize(
      attn_mask: torch.Tensor = None,
      size: Tuple[int] = None,
      device=get_current_device(),
  ):
      if attn_mask is not None:
          indices = (
              torch.nonzero(attn_mask.flatten(), as_tuple=False)
              .flatten()
              .to(device)
          )
          seqlens = attn_mask.sum(dim=-1, dtype=torch.int32).flatten()
      else:
          batch_size, tgt_len = size[0], size[1]
          indices = torch.arange(
              batch_size * tgt_len, dtype=torch.long, device=device
          )
          seqlens = torch.LongTensor(
              [tgt_len] * batch_size, device=device
          )
      max_seqlen = max(seqlens)
      cu_seqlens = F.pad(
          torch.cumsum(seqlens, dim=0, dtype=torch.int32), (1, 0)
      ).to(device)
      return SeqLenInfo(
          seqlens.tolist(), indices, max_seqlen, cu_seqlens
      )
  
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 SeqLenInfo obj's underlying member method code with example usage: obj.materialize(attn_mask, size, device) and returns: SeqLenInfo