diff --git a/configs/experiment/prediction.yaml b/configs/experiment/prediction.yaml
new file mode 100644
index 0000000..738b1c0
--- /dev/null
+++ b/configs/experiment/prediction.yaml
@@ -0,0 +1,47 @@
+# @package _global_
+
+# to execute this experiment run:
+# python train.py experiment=example
+
+defaults:
+  - override /model: s2bms_prediction
+  - override /data: s2bms_prediction
+  - override /metrics: s2bms_predictive
+
+# all parameters below will be merged with parameters from default configurations set above
+# this allows you to overwrite only specified parameters
+
+model:
+  geo_encoder:
+    _target_: src.models.components.geo_encoders.geoclip.GeoClipCoordinateEncoder
+    geo_data_name: coords
+  trainable_modules: [prediction_head]
+  prediction_head:
+    _target_: src.models.components.pred_heads.linear_pred_head.LinearPredictionHead
+
+data:
+  dataset:
+    modalities:
+      coords:
+    use_target_data: true
+    use_aux_data: false
+  caption_builder:
+  pin_memory: false
+  batch_size: 8
+  persistent_workers: false
+  saved_split_file_name: "s2bms_union_val_test.pth"
+
+tags: ["prediction", "geoclip_coords"]
+
+seed: 12345
+
+trainer:
+  min_epochs: 1
+  max_epochs: 100
+
+logger:
+  wandb:
+    tags: ${tags}
+    group: "predictive"
+  aim:
+    experiment: "predictive"
diff --git a/data/s2bms/splits/explanations.txt b/data/s2bms/splits/explanations.txt
new file mode 100644
index 0000000..56c5134
--- /dev/null
+++ b/data/s2bms/splits/explanations.txt
@@ -0,0 +1,6 @@
+split_indices_s2bms-unlabelled-20260529_2026-05-29-1433.pth: all train/val/test unlabelled indices (not used, source data)
+split_indices_s2bms+s2bms-unlabelled-20260529_2026-05-29-1438.pth: original S2BMS split and all train/val/test unlabelled indices (not used, source data)
+split_indices_s2bms_2024-08-14-1459.pth: original S2BMS split (not used, source data)
+
+s2bms_union_val_test: S2BMS labelled subset for this study. All val and test split datapoints are available for modalities of: coords, alphaearth and tessera.
+s2bms_unlabelled_union_val_test: s2bms_union_val_test + all train samples from the unlabelled dataset
\ No newline at end of file
diff --git a/data/s2bms/splits/split_indices_s2bms_2024-08-14-1459_union_aef_tsr_val_test.pth b/data/s2bms/splits/s2bms_union_val_test.pth
similarity index 100%
rename from data/s2bms/splits/split_indices_s2bms_2024-08-14-1459_union_aef_tsr_val_test.pth
rename to data/s2bms/splits/s2bms_union_val_test.pth
diff --git a/data/s2bms/splits/s2bms_unlabelled_union_val_test.pth b/data/s2bms/splits/s2bms_unlabelled_union_val_test.pth
new file mode 100644
index 0000000..fa57dc0
Binary files /dev/null and b/data/s2bms/splits/s2bms_unlabelled_union_val_test.pth differ
diff --git a/data/s2bms/splits/split_indices_s2bms_2024-08-14-1459_union_aef_tsr_val_test_train.pth b/data/s2bms/splits/split_indices_s2bms_2024-08-14-1459_union_aef_tsr_val_test_train.pth
deleted file mode 100644
index aed7e7d..0000000
Binary files a/data/s2bms/splits/split_indices_s2bms_2024-08-14-1459_union_aef_tsr_val_test_train.pth and /dev/null differ
diff --git a/notebooks/09-GT-aef-tessera-datasplit-filtering.ipynb b/notebooks/09-GT-aef-tessera-datasplit-filtering.ipynb
index 6391757..020dcf6 100644
--- a/notebooks/09-GT-aef-tessera-datasplit-filtering.ipynb
+++ b/notebooks/09-GT-aef-tessera-datasplit-filtering.ipynb
@@ -25,13 +25,11 @@
    "outputs": [],
    "source": [
     "# Check how do aef and tessera files overlap\n",
-    "\n",
-    "paths = glob.glob(\"data/s2bms/eo/aef/*.tif\")\n",
+    "paths = glob.glob(\"data/s2bms/eo/aef/*.npy\")\n",
     "aef = []\n",
     "for p in paths:\n",
     "    aef.append(os.path.basename(p).split(\".\")[0].split(\"-\")[-1])\n",
     "\n",
-    "\n",
     "paths = glob.glob(\"data/s2bms/eo/tessera/*.npy\")\n",
     "tsr = []\n",
     "for p in paths:\n",
@@ -42,6 +40,10 @@
     "tsr_left = list(set(tsr) - set(common))  # None\n",
     "aef_left = list(set(aef) - set(common))  # 12\n",
     "\n",
+    "# Filter out unlabelled data\n",
+    "tsr_left = {i for i in tsr_left if \"sample\" not in i}\n",
+    "aef_left = {i for i in aef_left if \"sample\" not in i}\n",
+    "\n",
     "print(\"Tessera:\", tsr_left)\n",
     "print(\"AlphaEarth:\", aef_left)"
    ]
@@ -98,14 +100,12 @@
     "        print(k, len(v))\n",
     "        s += len(v)\n",
     "\n",
-    "if os.path.exists(\n",
-    "    \"data/s2bms/splits/split_indices_s2bms_2024-08-14-1459_union_aef_tsr_val_test.pth\"\n",
-    "):\n",
+    "if os.path.exists(\"data/s2bms/splits/s2bms_union_val_test.pth\"):\n",
     "    print(\"Already saved\")\n",
     "else:\n",
     "    torch.save(\n",
     "        split_indices,\n",
-    "        \"data/s2bms/splits/split_indices_s2bms_2024-08-14-1459_union_aef_tsr_val_test.pth\",\n",
+    "        \"data/s2bms/splits/s2bms_union_val_test.pth\",\n",
     "    )"
    ]
   },
@@ -116,59 +116,182 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# Filter out remaining aef but missing tessera tiles form train splits\n",
+    "# Visualise data split\n",
+    "df = pd.read_csv(\"data/s2bms/model_ready_s2bms.csv\")\n",
     "\n",
-    "split_indices[\"train_indices\"] = split_indices[\"train_indices\"][\n",
-    "    ~split_indices[\"train_indices\"].isin(aef_left_str)\n",
-    "]\n",
+    "train_idx = split_indices[\"train_indices\"]\n",
+    "train_df = df[df[\"name_loc\"].isin(map(str, train_idx))]\n",
+    "\n",
+    "val_idx = split_indices[\"val_indices\"]\n",
+    "val_df = df[df[\"name_loc\"].isin(map(str, val_idx))]\n",
+    "\n",
+    "test_idx = split_indices[\"test_indices\"]\n",
+    "test_df = df[df[\"name_loc\"].isin(map(str, test_idx))]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "6",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import folium\n",
+    "\n",
+    "m = folium.Map(location=[df.lat.mean(), df.lon.mean()], zoom_start=3)\n",
+    "\n",
+    "# Train = blue\n",
+    "for _, row in train_df.iterrows():\n",
+    "    folium.CircleMarker(location=[row.lat, row.lon], radius=2, color=\"blue\", fill=True).add_to(m)\n",
+    "\n",
+    "# Val = orange\n",
+    "for _, row in val_df.iterrows():\n",
+    "    folium.CircleMarker(location=[row.lat, row.lon], radius=2, color=\"orange\", fill=True).add_to(m)\n",
+    "\n",
+    "# Test = red\n",
+    "for _, row in test_df.iterrows():\n",
+    "    folium.CircleMarker(location=[row.lat, row.lon], radius=2, color=\"red\", fill=True).add_to(m)\n",
+    "\n",
+    "m"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "7",
+   "metadata": {},
+   "source": [
+    "# Re-do Unlabelled split"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "8",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "df_unlabelled = pd.read_csv(\"data/s2bms/model_ready_s2bms-unlabelled-20260529.csv\")\n",
+    "df_unlabelled = df_unlabelled[df_unlabelled.split == \"train\"]\n",
+    "\n",
+    "train_unlabelled = list(df_unlabelled.name_loc)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "9",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "pth = \"data/s2bms/splits/split_indices_s2bms+s2bms-unlabelled-20260529_2026-05-29-1438.pth\"\n",
+    "split_indices = torch.load(pth, weights_only=False)\n",
     "\n",
-    "print(\"After test/val/train cleaning:\")\n",
+    "print(\"Before test/val/train:\")\n",
+    "s = 0\n",
     "for k, v in split_indices.items():\n",
     "    if not k == \"clusters\":\n",
     "        print(k, len(v))\n",
-    "        s += len(v)\n",
-    "\n",
-    "if os.path.exists(\n",
-    "    \"data/s2bms/splits/split_indices_s2bms_2024-08-14-1459_union_aef_tsr_val_test_train.pth\"\n",
-    "):\n",
-    "    print(\"Already saved\")\n",
-    "else:\n",
+    "        s += len(v)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "10",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "split_indices[\"train_indices\"] = pd.Series(\n",
+    "    [i for i in list(split_indices[\"train_indices\"]) if i in train_unlabelled or \"UKBMS\" in i]\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "11",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "val_test_indices = torch.load(\n",
+    "    \"data/s2bms/splits/s2bms_union_val_test.pth\",\n",
+    "    weights_only=False,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "12",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "split_indices[\"val_indices\"] = pd.Series(\n",
+    "    [i for i in list(val_test_indices[\"val_indices\"]) if \"UKBMS\" in i]\n",
+    ")\n",
+    "split_indices[\"test_indices\"] = pd.Series(\n",
+    "    [i for i in list(val_test_indices[\"test_indices\"]) if \"UKBMS\" in i]\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "13",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "print(\"After test/val/train:\")\n",
+    "s = 0\n",
+    "for k, v in split_indices.items():\n",
+    "    if not k == \"clusters\":\n",
+    "        print(k, len(v))\n",
+    "        s += len(v)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "14",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "if not os.path.exists(\"data/s2bms/splits/s2bms_unlabelled_union_val_test.pth\"):\n",
     "    torch.save(\n",
     "        split_indices,\n",
-    "        \"data/s2bms/splits/split_indices_s2bms_2024-08-14-1459_union_aef_tsr_val_test_train.pth\",\n",
-    "    )"
+    "        \"data/s2bms/splits/s2bms_unlabelled_union_val_test.pth\",\n",
+    "    )\n",
+    "else:\n",
+    "    print(\"Already saved\")"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "6",
+   "id": "15",
    "metadata": {},
    "outputs": [],
    "source": [
     "# Visualise data split\n",
-    "df = pd.read_csv(\"data/s2bms/model_ready_s2bms.csv\")\n",
+    "df_merged = pd.concat([df, df_unlabelled])\n",
     "\n",
     "train_idx = split_indices[\"train_indices\"]\n",
-    "train_df = df[df[\"name_loc\"].isin(map(str, train_idx))]\n",
-    "\n",
+    "train_df = df_merged[df_merged[\"name_loc\"].isin(map(str, train_idx))]\n",
     "\n",
     "val_idx = split_indices[\"val_indices\"]\n",
-    "val_df = df[df[\"name_loc\"].isin(map(str, val_idx))]\n",
+    "val_df = df_merged[df_merged[\"name_loc\"].isin(map(str, val_idx))]\n",
     "\n",
     "test_idx = split_indices[\"test_indices\"]\n",
-    "test_df = df[df[\"name_loc\"].isin(map(str, test_idx))]"
+    "test_df = df_merged[df_merged[\"name_loc\"].isin(map(str, test_idx))]"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "7",
+   "id": "16",
    "metadata": {},
    "outputs": [],
    "source": [
-    "import folium\n",
-    "\n",
     "m = folium.Map(location=[df.lat.mean(), df.lon.mean()], zoom_start=3)\n",
     "\n",
     "# Train = blue\n",
@@ -189,13 +312,19 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "8",
+   "id": "17",
    "metadata": {},
    "outputs": [],
    "source": []
   }
  ],
- "metadata": {},
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  }
+ },
  "nbformat": 4,
  "nbformat_minor": 5
 }
diff --git a/src/data/base_datamodule.py b/src/data/base_datamodule.py
index 1793d26..7f969cf 100644
--- a/src/data/base_datamodule.py
+++ b/src/data/base_datamodule.py
@@ -2,7 +2,7 @@
 import os
 import time
 from functools import partial
-from typing import Any, Dict, List, Tuple
+from typing import Any, Tuple
 
 import numpy as np
 import pandas as pd
@@ -67,8 +67,6 @@ def __init__(
             self.caption_builder.sync_with_dataset(self.dataset)
             self.concept_configs = caption_builder.concepts
 
-        self.split_data()
-
     @property
     def tabular_dim(self):
         return self.dataset.tabular_dim
@@ -86,6 +84,7 @@ def setup(self, stage: str = "fit") -> None:
 
         # Set up the dataset (download requested modalities)
         self.dataset.setup()
+        self.split_data()
 
     @property
     def batch_size_per_device(self) -> None:
@@ -251,10 +250,19 @@ def split_data(self) -> None:
             if test_indices is not None and not isinstance(test_indices, pd.Series):
                 raise NotImplementedError("Expected a pd series of name_locs for data splits.")
 
-            train_indices = np.where(self.dataset.df["name_loc"].isin(train_indices))[0]
-            val_indices = np.where(self.dataset.df["name_loc"].isin(val_indices))[0]
+            ds_records_names = [i["name_loc"] for i in self.dataset.records]
+            records_name_to_idx = {name: idx for idx, name in enumerate(ds_records_names)}
+
+            train_indices = np.array(
+                [records_name_to_idx[n] for n in train_indices if n in records_name_to_idx]
+            )
+            val_indices = np.array(
+                [records_name_to_idx[n] for n in val_indices if n in records_name_to_idx]
+            )
             if test_indices is not None:
-                test_indices = np.where(self.dataset.df["name_loc"].isin(test_indices))[0]
+                test_indices = np.array(
+                    [records_name_to_idx[n] for n in test_indices if n in records_name_to_idx]
+                )
 
             print(f"Dataset was split using indices from file: {self.saved_split_file_path}")
         else:
@@ -264,13 +272,10 @@ def split_data(self) -> None:
 
         if split_data_from_inds:
             self.data_train = torch.utils.data.Subset(self.dataset, train_indices)
-            self.data_train.dataset.mode = "train"
             self.data_val = torch.utils.data.Subset(self.dataset, val_indices)
-            self.data_val.dataset.mode = "val"
 
             if test_indices is not None:
                 self.data_test = torch.utils.data.Subset(self.dataset, test_indices)
-                self.data_test.dataset.mode = "test"
             else:
                 self.data_test = None
 
diff --git a/src/data/base_dataset.py b/src/data/base_dataset.py
index d39e66a..64ccd7b 100644
--- a/src/data/base_dataset.py
+++ b/src/data/base_dataset.py
@@ -10,7 +10,6 @@
 from torch.utils.data import Dataset
 
 from src.utils.data_utils import center_crop_npy
-from src.utils.errors import MissingDataError
 
 TORCH_DTYPES = {
     "float32": torch.float32,
@@ -262,19 +261,22 @@ def setup_tessera(self) -> None:
 
         # If data does not exist or is empty → full download
         if not os.path.exists(dst_dir) or len(os.listdir(dst_dir)) == 0:
-            os.makedirs(dst_dir, exist_ok=True)
-
-            tessera_from_df(
-                self.df,
-                data_dir=dst_dir,
-                year=year,
-                tile_size=size,
-                cache_dir=self.cache_dir,
-            )
+            if download_missing_tiles:
+                os.makedirs(dst_dir, exist_ok=True)
+
+                tessera_from_df(
+                    self.df,
+                    data_dir=dst_dir,
+                    year=year,
+                    tile_size=size,
+                    cache_dir=self.cache_dir,
+                )
 
-            # TODO: if we compile the dataset and use zenodo (or sth else) then change to pooch downloading/loading
-            # TODO: in case of zenodo use may need to be moved to UC dataset subclasses
-            # or self.setup_tessera_from_pooch() <- per children class implementation
+                # TODO: if we compile the dataset and use zenodo (or sth else) then change to pooch downloading/loading
+                # TODO: in case of zenodo use may need to be moved to UC dataset subclasses
+                # or self.setup_tessera_from_pooch() <- per children class implementation
+            else:
+                print("Please download the missing Tessera tiles...")
 
         # Download missing rows (if any)
         else:
@@ -306,8 +308,9 @@ def setup_tessera(self) -> None:
                             continue
                         except Exception as e:
                             print(f"Tile for {fname} could not be retrieved. Error: {e}")
-                self.records.pop(i)
-                print(f"No tile found for {fname} thus it will not be used.")
+                    else:
+                        self.records.pop(i)
+                        print(f"No tile found for {fname} thus it will not be used.")
 
     @final
     def setup_aef(self) -> None:
@@ -319,6 +322,12 @@ def setup_aef(self) -> None:
         print("\n\nSetting up AEF data...\n\n")
 
         dst_dir = os.path.join(self.data_dir, "eo/aef")
+        avail_files = os.listdir(dst_dir)
+        for i, rec in enumerate(self.records):
+            fname = os.path.basename(rec["aef_path"])
+            if fname not in avail_files:
+                self.records.pop(i)
+                print(f"No tile found for {fname} thus it will not be used.")
 
         # TODO aef retrieval?
         # TODO: in case of zenodo use may need to be moved to UC dataset subclasses
diff --git a/tests/test_captions.py b/tests/test_captions.py
index 85951dd..454f533 100644
--- a/tests/test_captions.py
+++ b/tests/test_captions.py
@@ -49,6 +49,7 @@ def test_datamodule_uses_collate_when_aux_data(request, sample_csv, tmp_path):
         num_workers=0,
         pin_memory=False,
     )
+    dm.setup()
 
     batch = next(iter(dm.train_dataloader()))
     assert "text" in batch
diff --git a/tests/test_datasets_and_datamodules.py b/tests/test_datasets_and_datamodules.py
index 20c94be..acdae4a 100644
--- a/tests/test_datasets_and_datamodules.py
+++ b/tests/test_datasets_and_datamodules.py
@@ -1,4 +1,3 @@
-from src.data.base_dataset import BaseDataset
 from src.data.butterfly_dataset import ButterflyDataset
 from src.data.heat_guatemala_dataset import HeatGuatemalaDataset
 from src.data.satbird_dataset import SatBirdDataset
@@ -64,6 +63,7 @@ def test_datasets_generic_properties(request, tmp_path, sample_csv):
 
 def test_datamodule_random_split_and_loaders(create_butterfly_dataset):
     dataset, dm = create_butterfly_dataset
+    dm.setup()
 
     assert len(dm.data_train) == 4
     assert len(dm.data_val) == 1
@@ -77,6 +77,7 @@ def test_datamodule_random_split_and_loaders(create_butterfly_dataset):
 def test_random_split_is_deterministic(create_butterfly_dataset):
     dataset1, dm1 = create_butterfly_dataset
     dataset2, dm2 = create_butterfly_dataset
+    dm1.setup(), dm2.setup()
 
     assert dm1.data_train.indices == dm2.data_train.indices
     assert dm1.data_val.indices == dm2.data_val.indices
diff --git a/tests/test_geo_encoders.py b/tests/test_geo_encoders.py
index 57ff921..4f062fb 100644
--- a/tests/test_geo_encoders.py
+++ b/tests/test_geo_encoders.py
@@ -1,8 +1,3 @@
-import json
-import os
-
-import pandas as pd
-import pytest
 import torch
 
 from src.models.components.geo_encoders.average_encoder import AverageEncoder
@@ -23,6 +18,7 @@ def test_geo_encoder_generic_properties(create_butterfly_dataset):
         "mlp_projector": MLPProjector,
     }
     ds, dm = create_butterfly_dataset
+    dm.setup()
     batch = next(iter(dm.train_dataloader()))
 
     for geo_encoder_name, geo_encoder_class in dict_geo_encoders.items():
diff --git a/tests/test_pred_heads.py b/tests/test_pred_heads.py
index 6f617c2..36149e0 100644
--- a/tests/test_pred_heads.py
+++ b/tests/test_pred_heads.py
@@ -1,12 +1,6 @@
-import json
-import os
-
-import pandas as pd
-import pytest
 import torch
 
 from src.models.components.geo_encoders.geoclip import GeoClipCoordinateEncoder
-from src.models.components.pred_heads.base_pred_head import BasePredictionHead
 from src.models.components.pred_heads.linear_pred_head import LinearPredictionHead
 from src.models.components.pred_heads.mlp_pred_head import MLPPredictionHead
 from src.models.components.pred_heads.mlp_regression_head import (
@@ -18,6 +12,7 @@
 def test_pred_head_generic_properties(create_butterfly_dataset):
     """Test required properties of the prediction head class."""
     ds, dm = create_butterfly_dataset
+    dm.setup()
     batch = next(iter(dm.train_dataloader()))
     eo_encoder = GeoClipCoordinateEncoder()
     eo_encoder.setup()
diff --git a/tests/test_text_encoders.py b/tests/test_text_encoders.py
index 0d09714..6ec31fa 100644
--- a/tests/test_text_encoders.py
+++ b/tests/test_text_encoders.py
@@ -1,11 +1,6 @@
-import json
-import os
-
-import pandas as pd
 import pytest
 import torch
 
-from src.models.components.text_encoders.base_text_encoder import BaseTextEncoder
 from src.models.components.text_encoders.clip_text_encoder import ClipTextEncoder
 from src.models.components.text_encoders.llm2clip_text_encoder import (
     LLM2CLIPTextEncoder,
@@ -18,6 +13,7 @@ def test_text_encoder_generic_properties(create_butterfly_dataset):
     """This test checks that all text encoders implement the basic properties and methods."""
     list_text_encoders = [ClipTextEncoder, LLM2CLIPTextEncoder]
     ds, dm = create_butterfly_dataset
+    dm.setup()
     batch = next(iter(dm.train_dataloader()))
     text_input = batch.get("text")
     assert text_input is not None, f"text input is None in the batch from datamodule. {batch}"
diff --git a/tests/test_train.py b/tests/test_train.py
index a60a1e6..52b5eca 100644
--- a/tests/test_train.py
+++ b/tests/test_train.py
@@ -88,6 +88,8 @@ def test_train_resume(tmp_path: Path, cfg_train: DictConfig) -> None:
     """
     with open_dict(cfg_train):
         cfg_train.trainer.max_epochs = 1
+        cfg_train.callbacks.model_checkpoint.save_last = True
+        cfg_train.callbacks.model_checkpoint.dirpath = cfg_train.paths.output_dir + "/checkpoints"
 
     HydraConfig().set_config(cfg_train)
     metric_dict_1, _ = train(cfg_train)
diff --git a/tests/test_yield_africa.py b/tests/test_yield_africa.py
index 69cc138..d311c7d 100644
--- a/tests/test_yield_africa.py
+++ b/tests/test_yield_africa.py
@@ -282,6 +282,7 @@ def test_yield_africa_dataset_aux_shape(yield_africa_dataset_with_aux):
 
 def test_yield_africa_datamodule_split_sizes(yield_africa_datamodule):
     dm = yield_africa_datamodule
+    dm.setup()
     assert len(dm.data_train) == 7
     assert len(dm.data_val) == 2
     assert len(dm.data_test) == 1
@@ -323,6 +324,7 @@ def make_dm():
         )
 
     dm1, dm2 = make_dm(), make_dm()
+    dm1.setup(), dm2.setup()
     assert dm1.data_train.indices == dm2.data_train.indices
     assert dm1.data_val.indices == dm2.data_val.indices