Repeating tensor different number of times fully on GPU

I am trying to implement a model in which there are categories (genera) and subcategories (species) which I am trying to classify images into. I want to build a model which explicitly separates the output into P[species|data] = P[genus | data] * P[species | genus, data]. I am implementing this by having one layer compute Log[P[genus|data]] and adding that to another layer which computes Log[P[species | genus, data]]. Unfortunately the various genera have different numbers of species in the dataset. So, in order to compute the final output I need to transform Log[P[genus|data]] into a tensor where each entry is repeated a number of times equal to the number of species in each genus. I am achieving this right now in a very inefficient way which relies on a python list comprehension as follows: (G is a tensor of size [batch_size, n_genera] giving the genus log-probabilities and S is a tensor of size [batch_size, n_species] giving the species log-probabilities with n_species > n_genera):

genus_expand = torch.transpose(torch.stack([g for s,g in zip(n_species_list,torch.transpose(G,0,1)) for i in range(s)],0),0,1)

final_prob = S + G

where n_species_list is a list of length n_genera, of integers giving the number of species in each genera.

Anyway, all that is just to explain what I am doing. Implemented this way I only seem to be able to get about 70% GPU utilization (an implementation without this line, where I only use P[species|data] so its entirely on the GPU and no python list comprehension) gets more like 90-95% utilization). I want to use something like tile but since the number of repetitions is different for different genera I don’t think I can make it work without essentially reproducing the above line of code.

Let me know if my explanation of my goal here is not fully clear.

It just occurred to me that I can solve this by constructing a matrix on the GPU to convert G to its expanded version. This seems like a good solution, though it seems rather inelegant… if anyone has a better idea let me know!