Added NotImplementedError to remember testing EnsembleFeedForwardNN

datamol-io · Dec 13, 2023 · d926c28 · d926c28
1 parent dd52ca5
commit d926c28
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 0 deletions.
diff --git a/graphium/nn/architectures/__init__.py b/graphium/nn/architectures/__init__.py
@@ -3,3 +3,4 @@
 from .global_architectures import TaskHeads
 from .global_architectures import GraphOutputNN
 from .pyg_architectures import FeedForwardPyg
+from .global_architectures import EnsembleFeedForwardNN
diff --git a/tests/test_ensemble_layers.py b/tests/test_ensemble_layers.py
@@ -14,6 +14,7 @@
     EnsembleMLP,
     EnsembleMuReadoutGraphium,
 )
+from graphium.nn.architectures import FeedForwardNN, EnsembleFeedForwardNN
 
 
 class test_Ensemble_Layers(ut.TestCase):
@@ -315,6 +316,8 @@ def test_ensemble_mlp(self):
             in_dim=11, out_dim=5, num_ensemble=3, batch_size=13, more_batch_dim=7, last_layer_is_readout=True
         )
 
+    def test_ensemble_feed_forward_nn(self):
+        raise NotImplementedError
 
 if __name__ == "__main__":
     ut.main()