diff --git a/neural_lam/datastore/base.py b/neural_lam/datastore/base.py
index 6859a82d..1be54ec2 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
         -------