Add safety check

PolicyEngine · Oct 24, 2024 · fa38af6 · fa38af6
1 parent c765176
commit fa38af6
Showing 1 changed file with 6 additions and 3 deletions.
diff --git a/policyengine_core/simulations/simulation.py b/policyengine_core/simulations/simulation.py
@@ -1551,9 +1551,12 @@ def subsample(
         self.build_from_dataset()
 
         # Ensure the baseline branch has the new data.
-        baseline_tax_benefit_system = self.baseline.tax_benefit_system
-        self.branches["baseline"] = self.clone()
-        self.branches["tax_benefit_system"] = baseline_tax_benefit_system
+        if "baseline" in self.branches:
+            baseline_tax_benefit_system = self.branches[
+                "baseline"
+            ].tax_benefit_system
+            self.branches["baseline"] = self.clone()
+            self.branches["tax_benefit_system"] = baseline_tax_benefit_system
         return self