Beam_search bottlenecks inference with only 1 used cpu

[adelplace]

It seems like I am not the only one facing this problem :

Model.generate() is extremely slow while using beam search

I was previously using trainer.predict() to run my inference but it does not offer that much flexibility and control over generation. For instance, I want to have more than one generation for each sample and it is not possible with trainer.predict(). So I changed to model.generate(). The problem is now I have to handle batching on my own and can not really use multiple GPUs. Still, I feel like it can not be that slow even under these circumstances. Below I have a small code snippet that does the…

Any ideas of solution ?