So if I understand this correctly, you have several variables and each of these has a moderate number of categories.

If they’re the same categories (i.e. you have the same embedding), it is more efficient to batch them.
If you have distinct embeddings, combining them would create a block-diagonal matrix with many zeros. That usually isn’t as efficient to process unless you use block-sparse variables, so I’d avoid that.

If, on the other hand, you have so many categories that using a linear layer is impractical, it is common to do negative sampling: If you first nn’s output is a vector probs, you can grab the K_pos largest indices (i_pos = topk(probs, K_pos, dim=-1).indices or somesuch) and K_neg random indices (i_neg = randint((batch, K_neg)); i_all = torch.cat((i_pos, i_neg), dim=-1)). Then you take probs[:, i_all] @ emb.weights.t()[i_all] + emb.bias() as embeddings.
This is not similar to negative sampling in word2vec (I recommend Richard Socher’s/Chris Manning’s lectures on them if you want a very detailed account), except that there the vectors are leaves and the network following them is shallow.
Best regards
Thomas