Repository navigation
Wrong return type in librispeech model_fn #762
Description
In librispeech_conformer the model_fn returns logits_batch as a Tuple of tensors, not a tensor.
The return type is hence wrong:
It should be:
def model_fn(...) -> Tuple[Tuple[spec.Tensor, spec.Tensor], spec.ModelAuxiliaryState]:
As insignificant as it seems, this caused me quite some trouble debugging an OOM issue. Might be useful for other ppl too.
Activity
mrhimanshu commentedon Oct 17, 2026
I reproduced this on the current dev branch. Both the PyTorch and JAX LibriSpeech Conformer implementations return a pair of tensors as the model output while annotating it as a single Tensor. I will update both annotations so the framework implementations stay consistent, verify the affected modules, and open a focused PR.