diff --git a/grain/python/__init__.py b/grain/python/__init__.py
index 0ac9d7d1f..c26e7cbb9 100644
--- a/grain/python/__init__.py
+++ b/grain/python/__init__.py
@@ -58,6 +58,7 @@
 from grain._src.python.dataset.base import (
     DatasetSelectionMap,
     RandomAccessDataSource,
+    SupportsBatchedReadRandomAccessDataSource,
 )
 from grain._src.python.dataset.dataset import (
     MapDataset,