Merge pull request graphnet-team#607 from AMHermansen/fix-tito

Fix tito
Aske-Rosted · Oct 24, 2023 · d2798fe · d2798fe
2 parents 8b9c353 + 111702d
commit d2798fe
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 4 deletions.
diff --git a/examples/04_training/02_train_tito_model.py b/examples/04_training/02_train_tito_model.py
@@ -182,8 +182,7 @@ def main(
 
     # Save model config and state dict - Version safe save method.
     model.save_state_dict(f"{path}/state_dict.pth")
-    # model.save_config(f"{path}/model_config.yml")
-    # Pending https://github.com/graphnet-team/graphnet/issues/606
+    model.save_config(f"{path}/model_config.yml")
 
 
 if __name__ == "__main__":

diff --git a/src/graphnet/models/gnn/dynedge_kaggle_tito.py b/src/graphnet/models/gnn/dynedge_kaggle_tito.py
@@ -35,7 +35,7 @@ class DynEdgeTITO(GNN):
     def __init__(
         self,
         nb_inputs: int,
-        features_subset: slice = slice(0, 4),
+        features_subset: List[int] = None,
         dyntrans_layer_sizes: Optional[List[Tuple[int, ...]]] = None,
         global_pooling_schemes: List[str] = ["max"],
     ):
@@ -120,7 +120,7 @@ def __init__(
         self._activation = torch.nn.LeakyReLU()
         self._nb_inputs = nb_inputs
         self._nb_global_variables = 5 + nb_inputs
-        self._features_subset = features_subset
+        self._features_subset = features_subset or [0, 1, 2, 3]
         self._construct_layers()
 
     def _construct_layers(self) -> None: