[Enhancement] Enable edge feature and edge input encoder for experime…

…ntal edage feature support (#920) *Issue #, if available:* *Description of changes:* This PR remove the two assertions of using edge features and edge input encoder to enable experimental edge feature support. By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice. Co-authored-by: Ubuntu <[email protected]>
awslabs · Jul 12, 2024 · 766d007 · 766d007
1 parent 1d0c41a
commit 766d007
Showing 1 changed file with 0 additions and 2 deletions.
diff --git a/python/graphstorm/model/gnn.py b/python/graphstorm/model/gnn.py
@@ -519,7 +519,6 @@ def get_dense_params(self):
             params += list(self.node_input_encoder.get_general_dense_parameters())
         # TODO(zhengda) we need to test a model with encoders on edge data.
         if self.edge_input_encoder is not None:
-            assert False, "edge_input_encoder not supported"
             params += list(self.edge_input_encoder.get_general_dense_parameters())
         if self.decoder is not None and isinstance(self.decoder, nn.Module):
             params += list(self.decoder.parameters())
@@ -538,7 +537,6 @@ def get_lm_params(self):
                 "node_input_encoder must be a GSNodeInputLayer"
             params += list(self.node_input_encoder.get_lm_dense_parameters())
         if self.edge_input_encoder is not None:
-            assert False, "edge_input_encoder not supported"
             params += list(self.edge_input_encoder.get_lm_dense_parameters())
 
         return params