Multiquery architectures broken - OpenNMT-py - score_batch

I keep running into a ValueError when using score_batch on a model architecture I've just converted. I'm using the score_batch function to filter some translation data to continue training the model with better data, but whenever the function is called on the Translator instance of the converted model checkpoint, it throws a shape error. The x and y seem to vary based on each individual batch, as when I try to just use score_batch on a single example there is no thrown error.

I've checked the structure/values of the passed in encoded source/target, and that all seems to be handled correctly. Passing in sentence batches of 4096-6048, with a max_batch_size of 2048 tokens. The processing of this function has worked for other architecture models with the exact code (no changes have been made). Translating using translate_batch works just fine, so I'm not sure what's going on.

Please let me know if a converted model / opennmt-py checkpoint would help or if I can assist.

Start by reproducing the reported shape error in score_batch with a converted OpenNMT-py checkpoint and multiquery architecture, comparing batched and single-example inputs. Compare this path with translate_batch, and use the batch-size details in the report to isolate the failing case; done means score_batch handles the converted model's batches without the ValueError.