diff --git a/neural_lam/datastore/base.py b/neural_lam/datastore/base.py
index 6859a82d..b8c7afa8 100644
--- a/neural_lam/datastore/base.py
+++ b/neural_lam/datastore/base.py
@@ -87,7 +87,7 @@ def get_vars_units(self, category: str) -> List[str]:
         Parameters
         ----------
         category : str
-            The category of the variables.
+            The category of the variables (state/forcing/static).
 
         Returns
         -------
@@ -104,7 +104,7 @@ def get_vars_names(self, category: str) -> List[str]:
         Parameters
         ----------
         category : str
-            The category of the variables.
+            The category of the variables (state/forcing/static).
 
         Returns
         -------
@@ -121,7 +121,7 @@ def get_num_data_vars(self, category: str) -> int:
         Parameters
         ----------
         category : str
-            The category of the variables.
+            The category of the variables (state/forcing/static).
 
         Returns
         -------