From 3464e3e2f66b9fbb3223b767695b5d246652d55d Mon Sep 17 00:00:00 2001
From: Eva Janouskova <janouskova.eva.012@gmail.com>
Date: Mon, 7 Oct 2024 23:00:52 +0100
Subject: [PATCH] lm: predict() does not always return pd.Series

---
 src/tlo/lm.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/src/tlo/lm.py b/src/tlo/lm.py
index e099714850..db0b10e716 100644
--- a/src/tlo/lm.py
+++ b/src/tlo/lm.py
@@ -456,6 +456,7 @@ def predict(
             # pop the boolean out of the series if we have a single row,
             # otherwise return the series
             if len(outcome) == 1 and squeeze_single_row_output:
+                assert isinstance(outcome.iloc[0], pd.Series)
                 return outcome.iloc[0]
             else:
                 return outcome