diff --git a/tensorflow_gnn/runner/tasks/regression.py b/tensorflow_gnn/runner/tasks/regression.py
index cb234ce4..977c0791 100644
--- a/tensorflow_gnn/runner/tasks/regression.py
+++ b/tensorflow_gnn/runner/tasks/regression.py
@@ -331,6 +331,7 @@ def __init__(self,
         node_set_name,
         units=units,
         state_name=state_name,
+        reduce_type=reduce_type,
         name=name,
         label_fn=label_fn,
         label_feature_name=label_feature_name)
@@ -369,6 +370,7 @@ def __init__(self,
         node_set_name,
         units=units,
         state_name=state_name,
+        reduce_type=reduce_type,
         name=name,
         label_fn=label_fn,
         label_feature_name=label_feature_name)
@@ -406,6 +408,7 @@ def __init__(self,
         node_set_name,
         units=units,
         state_name=state_name,
+        reduce_type=reduce_type,
         name=name,
         label_fn=label_fn,
         label_feature_name=label_feature_name)
@@ -444,6 +447,7 @@ def __init__(self,
         node_set_name,
         units=units,
         state_name=state_name,
+        reduce_type=reduce_type,
         name=name,
         label_fn=label_fn,
         label_feature_name=label_feature_name)
@@ -488,6 +492,7 @@ def __init__(self,
         node_set_name,
         units=units,
         state_name=state_name,
+        reduce_type=reduce_type,
         name=name,
         label_fn=label_fn,
         label_feature_name=label_feature_name,