diff --git a/usaon_benefit_tool/util/sankey.py b/usaon_benefit_tool/util/sankey.py
index 6143049..f728e88 100644
--- a/usaon_benefit_tool/util/sankey.py
+++ b/usaon_benefit_tool/util/sankey.py
@@ -180,7 +180,7 @@ def _node_ids_in_links(links: list[HighchartsSankeySeriesLink]) -> set[str]:
     return node_ids
 
 
-def _weight_for_criticality_rating(criticality_rating: int | None) -> float:
+def _weight_for_criticality_rating(criticality_rating: int | None) -> int | float:
     """If criticality rating is not set, return a very thin line."""
     if criticality_rating is None:
         return 0.01