diff --git a/tests/processor/test_batch_conversion.py b/tests/processor/test_batch_conversion.py
index 5cf8d3474..ae453ac3e 100644
--- a/tests/processor/test_batch_conversion.py
+++ b/tests/processor/test_batch_conversion.py
@@ -224,7 +224,7 @@ def test_complex_nested_observation():
 
     # Check that all observation keys are preserved
     original_obs_keys = {k for k in batch if k.startswith("observation.")}
-    reconstructed_obs_keys = {k for k in reconstructed_batch.keys() if k.startswith("observation.")}
+    reconstructed_obs_keys = {k for k in reconstructed_batch if k.startswith("observation.")}
 
     assert original_obs_keys == reconstructed_obs_keys