diff --git a/neural_lam/weather_dataset.py b/neural_lam/weather_dataset.py
index 5547fdd4..dfc00c78 100644
--- a/neural_lam/weather_dataset.py
+++ b/neural_lam/weather_dataset.py
@@ -607,10 +607,16 @@ def _is_listlike(obj):
 
         da_datastore_state = getattr(self, f"da_{category}")
         da_grid_index = da_datastore_state.grid_index
-        da_state_feature = da_datastore_state.state_feature
+        feature_key = f"{category}_feature"
+        if feature_key not in da_datastore_state.coords:
+            raise KeyError(
+                    f"Expected coordinate '{feature_key}' not found for category '{category}'. "
+                    f"Available coords: {list(da_datastore_state.coords.keys())}"
+            )
 
+        da_feature = da_datastore_state.coords[feature_key]
         coords = {
-            f"{category}_feature": da_state_feature,
+            feature_key: da_feature,
             "grid_index": da_grid_index,
         }
         if add_time_as_dim: