Fix passing through the reduce_type in graph regression tasks

PiperOrigin-RevId: 594226366
tensorflow · Dec 28, 2023 · 19c10f2 · 19c10f2
1 parent e6b78bd
commit 19c10f2
Showing 1 changed file with 5 additions and 0 deletions.
diff --git 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,