From b7ad65abddf5e2062aa0e2af5428893e78fba908 Mon Sep 17 00:00:00 2001
From: Himanshu <18554133+mrhimanshu@users.noreply.github.com>
Date: Sat, 17 Oct 2026 11:48:56 +0100
Subject: [PATCH] fix(librispeech): correct model_fn return type

---
 .../librispeech_jax/workload.py               |  2 +-
 .../librispeech_pytorch/workload.py           |  2 +-
 tests/test_model_fn_annotations.py            | 30 +++++++++++++++++++
 3 files changed, 32 insertions(+), 2 deletions(-)
 create mode 100644 tests/test_model_fn_annotations.py

diff --git a/algoperf/workloads/librispeech_conformer/librispeech_jax/workload.py b/algoperf/workloads/librispeech_conformer/librispeech_jax/workload.py
index 8eb37e25b..2b2b323a4 100644
--- a/algoperf/workloads/librispeech_conformer/librispeech_jax/workload.py
+++ b/algoperf/workloads/librispeech_conformer/librispeech_jax/workload.py
@@ -106,7 +106,7 @@ def model_fn(
     update_batch_norm: bool,
     use_running_average_bn: Optional[bool] = None,
     dropout_rate: Optional[float] = models.DROPOUT_RATE,
-  ) -> Tuple[spec.Tensor, spec.ModelAuxiliaryState]:
+  ) -> Tuple[Tuple[spec.Tensor, spec.Tensor], spec.ModelAuxiliaryState]:
     variables = {'params': params, **model_state}
     inputs, input_paddings = augmented_and_preprocessed_input_batch['inputs']
     is_train_mode = mode == spec.ForwardPassMode.TRAIN
diff --git a/algoperf/workloads/librispeech_conformer/librispeech_pytorch/workload.py b/algoperf/workloads/librispeech_conformer/librispeech_pytorch/workload.py
index 25416682c..7821f8357 100644
--- a/algoperf/workloads/librispeech_conformer/librispeech_pytorch/workload.py
+++ b/algoperf/workloads/librispeech_conformer/librispeech_pytorch/workload.py
@@ -104,7 +104,7 @@ def model_fn(
     rng: spec.RandomState,
     update_batch_norm: bool,
     dropout_rate: float = models.DROPOUT_RATE,
-  ) -> Tuple[spec.Tensor, spec.ModelAuxiliaryState]:
+  ) -> Tuple[Tuple[spec.Tensor, spec.Tensor], spec.ModelAuxiliaryState]:
     del model_state
     del rng
 
diff --git a/tests/test_model_fn_annotations.py b/tests/test_model_fn_annotations.py
new file mode 100644
index 000000000..269ff7bbd
--- /dev/null
+++ b/tests/test_model_fn_annotations.py
@@ -0,0 +1,30 @@
+import ast
+from pathlib import Path
+
+import pytest
+
+
+@pytest.mark.parametrize(
+  'workload_path',
+  [
+    Path(
+      'algoperf/workloads/librispeech_conformer/librispeech_jax/workload.py'
+    ),
+    Path(
+      'algoperf/workloads/librispeech_conformer/librispeech_pytorch/workload.py'
+    ),
+  ],
+)
+def test_librispeech_conformer_model_fn_return_annotation(
+  workload_path: Path,
+):
+  module = ast.parse(workload_path.read_text(encoding='utf-8'))
+  model_fn = next(
+    node
+    for node in ast.walk(module)
+    if isinstance(node, ast.FunctionDef) and node.name == 'model_fn'
+  )
+
+  assert ast.unparse(model_fn.returns) == (
+    'Tuple[Tuple[spec.Tensor, spec.Tensor], spec.ModelAuxiliaryState]'
+  )