I’m working with documents where given a query; I have to find the position of the “answer”. #queries varies across documents.
I have a document encoder and a query encoder. Normally, we could have
batch = {
    'doc':     [d1,  d1,  d2,  d2,  d2,  d2,  d3,  d4],
    'query':   [q11, q12, q21, q22, q23, q24, q31, q41],
    'y_pos':   [p11, p12, p21, p22, p23, p24, p31, p41],
}

As computing document embedding takes a while with RNN, I thought of doing:
batch = {
    'doc':     [d1,       d2,                 d3,  d4],
    'query':   [q11, q12, q21, q22, q23, q24, q31, q41],
    'y_pos':   [p11, p12, p21, p22, p23, p24, p31, p41],
    'nb_q':    [2,        4,                  1,   1  ],
}

So one model could compute each document embeddings with only 4 computations (d1, d2, d3,d4) and then use torch.repeat_interleave instead having 8 computations (2d1, 4d2, 1d3, 1d4).
Is it more clear ?
For your answer, I could split the model on 2 gpus but they aren’t big, both fit within a single gpu. Would be the gain worth ?