diff --git a/recml/core/utils/keras_utils.py b/recml/core/utils/keras_utils.py
index e271edc..edd0903 100644
--- a/recml/core/utils/keras_utils.py
+++ b/recml/core/utils/keras_utils.py
@@ -826,7 +826,7 @@ def _prepare_v3_restore(
   # Load index path metadata for non-trainable and optimizer variables if
   # needed.
   non_trainable_paths = None
-  if model is not None and abstract_state.get(NON_TRAINABLE_VARIABLES_KEY):
+  if model is not None and NON_TRAINABLE_VARIABLES_KEY in abstract_state:
     nt_path = epath.Path(checkpoint_path) / NON_TRAINABLE_PATHS_KEY
     checkpointer = ocp.Checkpointer(ocp.handlers.JsonCheckpointHandler())
     try:
@@ -840,7 +840,7 @@ def _prepare_v3_restore(
   if (
       model is not None
       and restore_optimizer_vars
-      and abstract_state.get(OPTIMIZER_VARIABLES_KEY)
+      and OPTIMIZER_VARIABLES_KEY in abstract_state
   ):
     opt_path = epath.Path(checkpoint_path) / OPTIMIZER_PATHS_KEY
     checkpointer = ocp.Checkpointer(ocp.handlers.JsonCheckpointHandler())
@@ -859,7 +859,7 @@ def _prepare_v3_restore(
   state_transforms = {}
 
   for key in keys:
-    if not abstract_state.get(key):
+    if key not in abstract_state:
       continue
 
     filtered_abstract_state[key] = {}
diff --git a/recml/core/utils/keras_utils_test.py b/recml/core/utils/keras_utils_test.py
index 88ae6bb..a0121bb 100644
--- a/recml/core/utils/keras_utils_test.py
+++ b/recml/core/utils/keras_utils_test.py
@@ -277,6 +277,38 @@ def test_keras_orbax_checkpointer_v3(
     # Ensures predictions are identical.
     np.testing.assert_allclose(preds, preds_after_restoration)
 
+  def test_keras_orbax_checkpointer_v3_no_non_trainable_variables(self):
+    model = keras.Sequential(
+        [keras.layers.Dense(5, name="dense")], name="model"
+    )
+    model.build((2, 10))
+    model.compile(optimizer=keras.optimizers.Adam(0.01), loss="mse")
+    model.optimizer.build(model.trainable_variables)
+    self.assertEmpty(model.non_trainable_variables)
+
+    checkpoint_dir = self.create_tempdir().full_path
+    checkpoint_manager = keras_utils.KerasOrbaxCheckpointManagerV3(
+        checkpoint_dir, max_to_keep=5
+    )
+    checkpoint_manager.save_model_variables(model, 0)
+    checkpoint_manager.wait_until_finished()
+
+    new_model = keras.Sequential(
+        [keras.layers.Dense(5, name="dense")], name="model"
+    )
+    new_model.build((2, 10))
+    new_model.compile(optimizer=keras.optimizers.Adam(0.01), loss="mse")
+    new_model.optimizer.build(new_model.trainable_variables)
+
+    checkpoint_manager.restore_model_variables(new_model, 0)
+    checkpoint_manager.close()
+
+    for expected, observed in zip(
+        [v.value for v in model.trainable_variables],
+        [v.value for v in new_model.trainable_variables],
+    ):
+      np.testing.assert_allclose(observed, expected)
+
   def test_restore_keras_checkpoint(self):
     dummy_inputs = _create_dummy_inputs()
     bert_pretrainer = _create_model(jax.tree.map(jnp.shape, dummy_inputs))