coerce population_size to be an integer

CDCgov · Feb 27, 2025 · abfdffb · abfdffb
1 parent 14ba0bb
commit abfdffb
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/pipelines/build_pyrenew_model.py b/pipelines/build_pyrenew_model.py
@@ -103,7 +103,7 @@ def build_model_from_dir(
         else None
     )
 
-    population_size = jnp.array(model_data["state_pop"])
+    population_size = jnp.array(model_data["state_pop"]).item()
 
     pop_fraction = jnp.array(model_data["pop_fraction"])
 

diff --git a/pyrenew_hew/pyrenew_wastewater_data.py b/pyrenew_hew/pyrenew_wastewater_data.py
@@ -51,7 +51,7 @@ def site_subpop_spine(self):
                         "site_pop": (
                             self.population_size
                             - site_indices.get_column("site_pop").sum()
-                        ).tolist(),
+                        ),
                     }
                 )
             else: