diff --git a/tests/integration/checkpointing_test.py b/tests/integration/checkpointing_test.py
index 5e263fa35f..134aadd5a2 100644
--- a/tests/integration/checkpointing_test.py
+++ b/tests/integration/checkpointing_test.py
@@ -232,5 +232,5 @@ def get_cmd(steps, metrics_file):
   # reports them as missing rather than leaving them at their init values.
   message = str(excinfo.value)
   assert "Checkpoint does not match the model" in message
-  assert "decoder/layers/0/" in message
+  assert "decoder/layers_0/" in message
   assert "scan_layers" in message