Improve test coverage.

aeonxai · Aug 16, 2023 · 83c0f7b · 83c0f7b
1 parent 0e769df
commit 83c0f7b
Showing 1 changed file with 13 additions and 0 deletions.
diff --git a/tests/layers/test_core_and_embedding.py b/tests/layers/test_core_and_embedding.py
@@ -84,3 +84,16 @@ def test_embedding(self, graph_Mo):
             node_attr, edge_attr, torch.tensor([0.0, 0.0])
         )  # this will be default value
         assert state_feat is None
+
+        # No ntypes_node.
+        embed5 = EmbeddingBlock(
+            degree_rbf=9,
+            dim_node_embedding=16,
+            dim_edge_embedding=16,
+            dim_state_feats=16,
+            include_state=True,
+            ntypes_node=None,
+            activation=nn.SiLU(),
+        )
+        node_feat, edge_feat, state_feat = embed5(node_attr, edge_attr, torch.tensor([1.0, 2.0]))
+        assert [state_feat.size(dim=0), state_feat.size(dim=1)] == [1, 16]