Is it possible to build batches of index?

[Wongboo]

Is it possible to build batches of index? Which is slightly different from batche defination in the doc, Each batch contains multiple queries and multiple databases, and queries only search for corresponding database in the batch. For example:

import torch

d = 64                           # dimension
batch_size=16              # batch size
nb = 4000                     # database size
nq = 2000                     # nb of queries
databases = torch.randn((batch_size, nb, d), device='cuda')
queries= torch.randn((batch_size, nq, d), device='cuda')
# doing some indexing and searchs, queries only search for corresponding database

I've searched whole wiki and discussions. Really appreciate your answer!