code: def _remained_idxs(mask: torch.Tensor, num_heads: int) -> List[int]:
      repeats = mask.shape[0] // num_heads
      remained = (
          mask.reshape([num_heads, -1]).sum(-1) != 0.0
      ).repeat_interleave(repeats)
      return torch.arange(len(mask))[remained].long().tolist()
  
file_path: nni/nni/compression/utils/external/external_replacer.py
imports: from __future__ import annotations
  from collections import defaultdict
  import logging
  import re
  from typing import TYPE_CHECKING, List
  import torch
  from torch.utils._pytree import tree_map
  from .huggingface import parser_factory, HuggingfaceModelParser
  from ..attr import get_nested_attr
  from ...speedup.replacer import Replacer
label_desc: Use image node_id 3 for calling a global function with example usage: _remained_idxs(mask, num_heads) and returns: torch