diff --git a/morph_net/network_regularizers/activation_regularizer.py b/morph_net/network_regularizers/activation_regularizer.py
index 6d6c4c6..427120b 100644
--- a/morph_net/network_regularizers/activation_regularizer.py
+++ b/morph_net/network_regularizers/activation_regularizer.py
@@ -68,7 +68,7 @@ def __init__(self,
     self._manager = orm.OpRegularizerManager(
         output_boundary,
         op_handler_dict,
-        input_boundary=input_boundary,
+        input_boundary=input_boundary,  # pyrefly: ignore[bad-argument-type]
         force_group=force_group,
         regularizer_blacklist=regularizer_blacklist)
     self._calculator = cost_calculator.CostCalculator(
@@ -153,7 +153,7 @@ def __init__(self,
     self._manager = orm.OpRegularizerManager(
         output_boundary,
         op_handler_dict,
-        input_boundary=input_boundary,
+        input_boundary=input_boundary,  # pyrefly: ignore[bad-argument-type]
         force_group=force_group,
         regularizer_blacklist=regularizer_blacklist)
     self._calculator = cost_calculator.CostCalculator(
diff --git a/morph_net/network_regularizers/flop_regularizer.py b/morph_net/network_regularizers/flop_regularizer.py
index 3f314f3..f970f36 100644
--- a/morph_net/network_regularizers/flop_regularizer.py
+++ b/morph_net/network_regularizers/flop_regularizer.py
@@ -84,7 +84,7 @@ def __init__(self,
     self._manager = orm.OpRegularizerManager(
         output_boundary,
         op_handler_dict,
-        input_boundary=input_boundary,
+        input_boundary=input_boundary,  # pyrefly: ignore[bad-argument-type]
         force_group=force_group,
         regularizer_blacklist=regularizer_blacklist)
     self._calculator = cost_calculator.CostCalculator(
@@ -157,7 +157,7 @@ def __init__(self,
     }
     if regularizer_decorator:
       for key in custom_handlers:
-        custom_handlers[key] = op_handler_decorator.OpHandlerDecorator(
+        custom_handlers[key] = op_handler_decorator.OpHandlerDecorator(  # pyrefly: ignore[unsupported-operation]
             custom_handlers[key], regularizer_decorator, decorator_parameters)
 
     op_handler_dict = op_handlers.get_group_lasso_op_handler_dict()
@@ -166,7 +166,7 @@ def __init__(self,
     self._manager = orm.OpRegularizerManager(
         output_boundary,
         op_handler_dict,
-        input_boundary=input_boundary,
+        input_boundary=input_boundary,  # pyrefly: ignore[bad-argument-type]
         force_group=force_group,
         regularizer_blacklist=regularizer_blacklist)
     self._calculator = cost_calculator.CostCalculator(
diff --git a/morph_net/network_regularizers/latency_regularizer.py b/morph_net/network_regularizers/latency_regularizer.py
index 5e4c0dd..0b92419 100644
--- a/morph_net/network_regularizers/latency_regularizer.py
+++ b/morph_net/network_regularizers/latency_regularizer.py
@@ -151,7 +151,7 @@ def __init__(self,
     })
 
     self._manager = orm.OpRegularizerManager(
-        output_boundary, op_handler_dict, input_boundary=input_boundary,
+        output_boundary, op_handler_dict, input_boundary=input_boundary,  # pyrefly: ignore[bad-argument-type]
         force_group=force_group, regularizer_blacklist=regularizer_blacklist)
     self._calculator = cost_calculator.CostCalculator(
         self._manager,
@@ -229,7 +229,7 @@ def __init__(self,
     }
     if regularizer_decorator:
       for key in custom_handlers:
-        custom_handlers[key] = op_handler_decorator.OpHandlerDecorator(
+        custom_handlers[key] = op_handler_decorator.OpHandlerDecorator(  # pyrefly: ignore[unsupported-operation]
             custom_handlers[key], regularizer_decorator, decorator_parameters)
 
     op_handler_dict = op_handlers.get_group_lasso_op_handler_dict()
@@ -238,7 +238,7 @@ def __init__(self,
     self._manager = orm.OpRegularizerManager(
         output_boundary,
         op_handler_dict,
-        input_boundary=input_boundary,
+        input_boundary=input_boundary,  # pyrefly: ignore[bad-argument-type]
         force_group=force_group,
         regularizer_blacklist=regularizer_blacklist)
     self._calculator = cost_calculator.CostCalculator(
diff --git a/morph_net/network_regularizers/logistic_sigmoid_regularizer.py b/morph_net/network_regularizers/logistic_sigmoid_regularizer.py
index e377fe8..3e5fe55 100644
--- a/morph_net/network_regularizers/logistic_sigmoid_regularizer.py
+++ b/morph_net/network_regularizers/logistic_sigmoid_regularizer.py
@@ -76,7 +76,7 @@ def __init__(self,
         output_boundary,
         op_handler_dict,
         create_grouping_regularizer=pgr.ProbabilisticGroupingRegularizer,
-        input_boundary=input_boundary,
+        input_boundary=input_boundary,  # pyrefly: ignore[bad-argument-type]
         force_group=force_group,
         regularizer_blacklist=regularizer_blacklist)
     self._calculator = self.get_calculator()
diff --git a/morph_net/network_regularizers/model_size_regularizer.py b/morph_net/network_regularizers/model_size_regularizer.py
index f245799..8877a6a 100644
--- a/morph_net/network_regularizers/model_size_regularizer.py
+++ b/morph_net/network_regularizers/model_size_regularizer.py
@@ -84,7 +84,7 @@ def __init__(self,
     })
 
     self._manager = orm.OpRegularizerManager(
-        output_boundary, op_handler_dict, input_boundary=input_boundary,
+        output_boundary, op_handler_dict, input_boundary=input_boundary,  # pyrefly: ignore[bad-argument-type]
         force_group=force_group, regularizer_blacklist=regularizer_blacklist)
     self._calculator = cost_calculator.CostCalculator(
         self._manager, resource_function.model_size_function)
@@ -157,7 +157,7 @@ def __init__(self,
     }
     if regularizer_decorator:
       for key in custom_handlers:
-        custom_handlers[key] = op_handler_decorator.OpHandlerDecorator(
+        custom_handlers[key] = op_handler_decorator.OpHandlerDecorator(  # pyrefly: ignore[unsupported-operation]
             custom_handlers[key], regularizer_decorator, decorator_parameters)
 
     op_handler_dict = op_handlers.get_group_lasso_op_handler_dict()
@@ -166,7 +166,7 @@ def __init__(self,
     self._manager = orm.OpRegularizerManager(
         output_boundary,
         op_handler_dict,
-        input_boundary=input_boundary,
+        input_boundary=input_boundary,  # pyrefly: ignore[bad-argument-type]
         force_group=force_group,
         regularizer_blacklist=regularizer_blacklist)
     self._calculator = cost_calculator.CostCalculator(