diff --git a/.github/workflows/pull_request.yaml b/.github/workflows/pull_request.yaml
index f25a2a8..c7992ad 100644
--- a/.github/workflows/pull_request.yaml
+++ b/.github/workflows/pull_request.yaml
@@ -49,6 +49,8 @@ jobs:
           HUGGING_FACE_TOKEN: ${{ secrets.HUGGING_FACE_TOKEN }}
       - name: Build datasets
         run: make data
+        env:
+          TEST_LITE: true
       - name: Run tests
         run: pytest
       - name: Test documentation builds
diff --git a/policyengine_uk_data/datasets/frs/local_areas/constituencies/calibrate.py b/policyengine_uk_data/datasets/frs/local_areas/constituencies/calibrate.py
index 1477e25..d824dd1 100644
--- a/policyengine_uk_data/datasets/frs/local_areas/constituencies/calibrate.py
+++ b/policyengine_uk_data/datasets/frs/local_areas/constituencies/calibrate.py
@@ -4,6 +4,7 @@
 import numpy as np
 from tqdm import tqdm
 import h5py
+import os
 from policyengine_uk_data.datasets.frs.local_areas.constituencies.transform_constituencies import (
     transform_2010_to_2024,
 )
@@ -69,7 +70,7 @@ def dropout_weights(weights, p):
 
     optimizer = torch.optim.Adam([weights], lr=0.05)
 
-    desc = range(2048)
+    desc = range(128) if os.environ.get("DATA_LITE") else range(2048)
 
     for epoch in desc:
         optimizer.zero_grad()
@@ -78,7 +79,7 @@ def dropout_weights(weights, p):
         l.backward()
         optimizer.step()
         if epoch % 50 == 0:
-            print(f"Loss: {l.item()}, Epoch: {epoch}")
+            print(f"Loss: {l.item()}, Epoch: {epoch}", flush=True)
 
     final_weights = torch.exp(weights).detach().numpy()
     mapping_matrix = pd.read_csv(
diff --git a/policyengine_uk_data/datasets/frs/local_areas/local_authorities/calibrate.py b/policyengine_uk_data/datasets/frs/local_areas/local_authorities/calibrate.py
index fa3bd3b..1dea9bc 100644
--- a/policyengine_uk_data/datasets/frs/local_areas/local_authorities/calibrate.py
+++ b/policyengine_uk_data/datasets/frs/local_areas/local_authorities/calibrate.py
@@ -4,10 +4,11 @@
 import numpy as np
 from tqdm import tqdm
 import h5py
+import os
 from policyengine_uk_data.storage import STORAGE_FOLDER
 
 
-from loss import (
+from policyengine_uk_data.datasets.frs.local_areas.local_authorities.loss import (
     create_local_authority_target_matrix,
     create_national_target_matrix,
 )
@@ -62,7 +63,7 @@ def dropout_weights(weights, p):
 
     optimizer = torch.optim.Adam([weights], lr=0.05)
 
-    desc = range(2048)
+    desc = range(128) if os.environ.get("DATA_LITE") else range(2048)
 
     for epoch in desc:
         optimizer.zero_grad()
diff --git a/policyengine_uk_data/utils/reweight.py b/policyengine_uk_data/utils/reweight.py
index 9f25d17..384168b 100644
--- a/policyengine_uk_data/utils/reweight.py
+++ b/policyengine_uk_data/utils/reweight.py
@@ -1,5 +1,6 @@
 import numpy as np
 import torch
+import os
 
 
 def reweight(
@@ -47,7 +48,7 @@ def dropout_weights(weights, p):
 
     start_loss = None
 
-    iterator = range(1_000)
+    iterator = range(128) if os.environ.get("DATA_LITE") else range(2048)
     for i in iterator:
         optimizer.zero_grad()
         weights_ = dropout_weights(weights, dropout_rate)