Wrong return type in librispeech model_fn

[Niccolo-Ajroldi]

In librispeech_conformer the model_fn returns logits_batch as a Tuple of tensors, not a tensor.

The return type is hence wrong:

algorithmic-efficiency/algorithmic_efficiency/workloads/librispeech_conformer/librispeech_pytorch/workload.py

Line 119 in ddf5efc

update_batch_norm: bool) -> Tuple[spec.Tensor, spec.ModelAuxiliaryState]:

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.