diff --git a/package-lock.json b/package-lock.json
index c852b36da..ca177a21a 100644
--- a/package-lock.json
+++ b/package-lock.json
@@ -20667,7 +20667,6 @@
       "resolved": "https://registry.npmjs.org/prettier/-/prettier-3.3.3.tgz",
       "integrity": "sha512-i2tDNA0O5IrMO757lfrdQZCc2jPNDVntV0m/+4whiDfWaTKfMNgR7Qz0NAeGz/nRqF4m5/6CLzbP4/liHt12Ew==",
       "dev": true,
-      "license": "MIT",
       "bin": {
         "prettier": "bin/prettier.cjs"
       },
diff --git a/src/data/reformDefinitionCode.js b/src/data/reformDefinitionCode.js
index 101f1abe6..96531633c 100644
--- a/src/data/reformDefinitionCode.js
+++ b/src/data/reformDefinitionCode.js
@@ -39,8 +39,11 @@ export function getHeaderCode(type, metadata, policy) {
     lines.push("from " + metadata.package + " import Microsimulation");
   }
 
-  // If there is a reform, add the following Python imports
-  if (Object.keys(policy.reform.data).length > 0) {
+  // If either baseline or reform is custom, add the following Python imports
+  if (
+    Object.keys(policy.reform.data).length > 0 ||
+    Object.keys(policy.baseline.data).length > 0
+  ) {
     lines.push("from policyengine_core.reforms import Reform");
   }