diff --git a/scripts/train_merge_unified.py b/scripts/train_merge_unified.py
index 99116c7..df515fa 100644
--- a/scripts/train_merge_unified.py
+++ b/scripts/train_merge_unified.py
@@ -100,15 +100,16 @@ def main():
     rng = random.Random(SEED)
     torch.manual_seed(SEED)
 
-    gsm = [dict(it, task="gsm") for it in json.load(open(OUT / "star_data.json"))
-           if it["split"] == "train" and it["label"] != "drop"]
-    mbpp = [dict(it, task="mbpp") for it in json.load(open(OUT / "mbpp_data.json"))
-            if it["split"] == "train" and it["label"] != "drop"
-            and it.get("sol_code")]
-    if "gsm" not in tasks:
-        gsm = []
-    if "mbpp" not in tasks:
-        mbpp = []
+    gsm, mbpp = [], []
+    if "gsm" in tasks:
+        gsm = [dict(it, task="gsm")
+               for it in json.load(open(OUT / "star_data.json"))
+               if it["split"] == "train" and it["label"] != "drop"]
+    if "mbpp" in tasks:
+        mbpp = [dict(it, task="mbpp")
+                for it in json.load(open(OUT / "mbpp_data.json"))
+                if it["split"] == "train" and it["label"] != "drop"
+                and it.get("sol_code")]
 
     model, tok = load_model(dtype=torch.bfloat16)
     for p in model.parameters():