diff --git a/examples/sequential_retrieval.py b/examples/sequential_retrieval.py
index c14555f..b72f741 100644
--- a/examples/sequential_retrieval.py
+++ b/examples/sequential_retrieval.py
@@ -56,6 +56,8 @@
 BATCH_SIZE = 2048
 TEST_BATCH_SIZE = 2048
 EMBEDDING_DIM = 128
+NUM_EPOCHS = 10
+LEARNING_RATE = 0.05
 
 """
 ## Dataset
@@ -372,8 +374,8 @@ def compute_loss(self, x, y, y_pred, sample_weight, training=True):
 
 # Compile.
 learning_rate = keras.optimizers.schedules.PolynomialDecay(
-    0.05,
-    decay_steps=10 * 30,
+    LEARNING_RATE,
+    decay_steps=train_ds.cardinality() * NUM_EPOCHS,
     end_learning_rate=0.0,
 )
 model.compile(optimizer=keras.optimizers.AdamW(learning_rate=learning_rate))
@@ -382,7 +384,7 @@ def compute_loss(self, x, y, y_pred, sample_weight, training=True):
 model.fit(
     train_ds,
     validation_data=test_ds,
-    epochs=30,
+    epochs=NUM_EPOCHS,
 )
 
 """