fix(tf2_gnn/models/*): let __init__ pass extra kwargs to super()

microsoft · Mar 11, 2021 · 16153fa · 16153fa
1 parent 9025691
commit 16153fa
Show file tree

Hide file tree

Showing 4 changed files with 6 additions and 9 deletions.
diff --git a/tf2_gnn/models/graph_binary_classification_task.py b/tf2_gnn/models/graph_binary_classification_task.py
@@ -18,9 +18,6 @@ def get_default_hyperparameters(
         super_params.update(these_hypers)
         return super_params
 
-    def __init__(self, params: Dict[str, Any], dataset: GraphDataset, name: str = None):
-        super().__init__(params, dataset=dataset, name=name)
-
     def compute_task_output(
         self,
         batch_features: Dict[str, tf.Tensor],

diff --git a/tf2_gnn/models/graph_regression_task.py b/tf2_gnn/models/graph_regression_task.py
@@ -30,8 +30,8 @@ def get_default_hyperparameters(
         super_params.update(these_hypers)
         return super_params
 
-    def __init__(self, params: Dict[str, Any], dataset: GraphDataset, name: str = None):
-        super().__init__(params, dataset=dataset, name=name)
+    def __init__(self, params: Dict[str, Any], dataset: GraphDataset, name: str = None, **kwargs):
+        super().__init__(params, dataset=dataset, name=name, **kwargs)
         self._node_to_graph_aggregation = None
 
         # Construct sublayers:

diff --git a/tf2_gnn/models/node_multiclass_task.py b/tf2_gnn/models/node_multiclass_task.py
@@ -31,8 +31,8 @@ def get_default_hyperparameters(cls, mp_style: Optional[str] = None) -> Dict[str
         super_params.update(these_hypers)
         return super_params
 
-    def __init__(self, params: Dict[str, Any], dataset: GraphDataset, name: str = None):
-        super().__init__(params, dataset=dataset, name=name)
+    def __init__(self, params: Dict[str, Any], dataset: GraphDataset, name: str = None, **kwargs):
+        super().__init__(params, dataset=dataset, name=name, **kwargs)
         if not hasattr(dataset, "num_node_target_labels"):
             raise ValueError(f"Provided dataset of type {type(dataset)} does not provide num_node_target_labels information.")
         self._num_labels = dataset.num_node_target_labels

diff --git a/tf2_gnn/models/qm9_regression.py b/tf2_gnn/models/qm9_regression.py
@@ -40,8 +40,8 @@ def get_default_hyperparameters(
         super_params.update(these_hypers)
         return super_params
 
-    def __init__(self, params: Dict[str, Any], dataset: GraphDataset, name: str = None):
-        super().__init__(params, dataset=dataset, name=name)
+    def __init__(self, params: Dict[str, Any], dataset: GraphDataset, name: str = None, **kwargs):
+        super().__init__(params, dataset=dataset, name=name, **kwargs)
         assert isinstance(dataset, QM9Dataset)
 
         self._task_id = int(dataset._params["task_id"])