code: def bmm(
      feat,
      nbrs,
      dist,
      sid,
      eid,
      sort=True,
      process_unit=4000,
      verbose=False,
  ):
      import torch
  
      _, cols = dist.shape
      batch_sim = torch.zeros((eid - sid, cols), dtype=torch.float32)
      for s in tqdm(
          range(sid, eid, process_unit), desc="bmm", disable=not verbose
      ):
          e = min(eid, s + process_unit)
          query = feat[s:e].unsqueeze(1)
          gallery = feat[nbrs[s:e]].permute(0, 2, 1)
          batch_sim[s - sid : e - sid] = torch.clamp(
              torch.bmm(query, gallery).view(-1, cols), 0.0, 1.0
          )
  
      if sort:
          sort_unit = int(1e6)
          batch_nbr = nbrs[sid:eid]
          for s in range(0, batch_sim.shape[0], sort_unit):
              e = min(s + sort_unit, eid)
              batch_sim[s:e], indices = torch.sort(
                  batch_sim[s:e], descending=True
              )
              batch_nbr[s:e] = torch.gather(batch_nbr[s:e], 1, indices)
          nbrs[sid:eid] = batch_nbr
      dist[sid:eid] = 1.0 - batch_sim
  
imports: import gc
  from tqdm import tqdm
  from .faiss_gpu import faiss_search_approx_knn
label_desc: Use image node_id 3 for calling a global function with example usage: bmm(feat, nbrs, dist, sid, eid, sort, process_unit, verbose) without return types