diff --git a/draco/core/containers.py b/draco/core/containers.py
index 14afe733f..0d8834bf9 100644
--- a/draco/core/containers.py
+++ b/draco/core/containers.py
@@ -480,7 +480,7 @@ def _make_selections(cls, sel_args):
 
         # Build selections dict
         selections = {}
-        for name, dataset in cls._dataset_spec.items():
+        for name, dataset in cls._class_dataset_spec().items():
             ds_axes = dataset["axes"]
             sel = []
             ds_relevant = False