diff --git a/.ipynb_checkpoints/main-checkpoint.ipynb b/.ipynb_checkpoints/main-checkpoint.ipynb
index 52390c0..bca1f60 100644
--- a/.ipynb_checkpoints/main-checkpoint.ipynb
+++ b/.ipynb_checkpoints/main-checkpoint.ipynb
@@ -9,6 +9,7 @@
    "source": [
     "from datetime import datetime, timedelta\n",
     "import pandas as pd\n",
+    "import numpy as np\n",
     "from tms_data_interface import SQLQueryInterface\n",
     "\n",
     "def apply_sar_flag(df, var1, var2, var3, random_state=42):\n",
diff --git a/main.ipynb b/main.ipynb
index 52390c0..bca1f60 100644
--- a/main.ipynb
+++ b/main.ipynb
@@ -9,6 +9,7 @@
    "source": [
     "from datetime import datetime, timedelta\n",
     "import pandas as pd\n",
+    "import numpy as np\n",
     "from tms_data_interface import SQLQueryInterface\n",
     "\n",
     "def apply_sar_flag(df, var1, var2, var3, random_state=42):\n",
diff --git a/main.py b/main.py
index af5c963..273fa8c 100644
--- a/main.py
+++ b/main.py
@@ -6,6 +6,7 @@
 
 from datetime import datetime, timedelta
 import pandas as pd
+import numpy as np
 from tms_data_interface import SQLQueryInterface
 
 def apply_sar_flag(df, var1, var2, var3, random_state=42):