diff --git a/algorithmic_efficiency/workloads/librispeech_deepspeech/librispeech_jax/workload.py b/algorithmic_efficiency/workloads/librispeech_deepspeech/librispeech_jax/workload.py
index 05fdf90e7..2d46960ed 100644
--- a/algorithmic_efficiency/workloads/librispeech_deepspeech/librispeech_jax/workload.py
+++ b/algorithmic_efficiency/workloads/librispeech_deepspeech/librispeech_jax/workload.py
@@ -1,5 +1,5 @@
 import functools
-from typing import Optional
+from typing import Optional, Dict, Tuple
 
 from flax import jax_utils
 import jax