diff --git a/.github/workflows/tests.yaml b/.github/workflows/tests.yaml
index b32a1bd..4a0dbf0 100644
--- a/.github/workflows/tests.yaml
+++ b/.github/workflows/tests.yaml
@@ -19,10 +19,10 @@ jobs:
       - name: Checkout
         uses: actions/checkout@v3
 
-      - name: Set up Python 3.11
+      - name: Set up Python 3.12
         uses: actions/setup-python@v3
         with:
-          python-version: "3.11"
+          python-version: "3.12"
 
       - name: Install packages
         run: |
diff --git a/src/MEDS_tabular_automl/tabularize.py b/src/MEDS_tabular_automl/tabularize.py
index 86948f5..a5ab4c8 100644
--- a/src/MEDS_tabular_automl/tabularize.py
+++ b/src/MEDS_tabular_automl/tabularize.py
@@ -39,7 +39,7 @@ def load_meds_data(MEDS_cohort_dir: str) -> Mapping[str, pl.DataFrame]:
     >>> assert "train" in split_to_df
     >>> assert len(split_to_df) == 3
     >>> assert len(split_to_df["train"]) == 1
-    >>> assert isinstance(split_to_df["train"][0], pl.DataFrame)
+    >>> assert isinstance(split_to_df["train"][0], pl.LazyFrame)
     """
     MEDS_cohort_dir = Path(MEDS_cohort_dir)
     meds_fps = list(MEDS_cohort_dir.glob("*/*.parquet"))
diff --git a/src/MEDS_tabular_automl/utils.py b/src/MEDS_tabular_automl/utils.py
index 0d703df..85ec597 100644
--- a/src/MEDS_tabular_automl/utils.py
+++ b/src/MEDS_tabular_automl/utils.py
@@ -211,7 +211,7 @@ def load_meds_data(MEDS_cohort_dir: str) -> Mapping[str, pl.DataFrame]:
     >>> assert "train" in split_to_df
     >>> assert len(split_to_df) == 3
     >>> assert len(split_to_df["train"]) == 1
-    >>> assert isinstance(split_to_df["train"][0], pl.DataFrame)
+    >>> assert isinstance(split_to_df["train"][0], pl.LazyFrame)
     """
     MEDS_cohort_dir = Path(MEDS_cohort_dir)
     meds_fps = list(MEDS_cohort_dir.glob("*/*.parquet"))