Beam Search without `model.generate`

I know there is model.generate to use beam search.

But in my case, I am using the BART encoder and decoder in a variational autoencoder.
As there are latent variables inside the autoencoder, I am not directly using the model.

In short,

enc = bart.encoder(input)
latent_var = some_linear_layers(enc)
output = bart.decoder(latent_var)

Is there any built-in way to use beam search in my case?