Fix pass of graph to graph layer

yura-hb · Mar 7, 2024 · d39a30a · d39a30a
1 parent 094808e
commit d39a30a
Showing 1 changed file with 1 addition and 5 deletions.
diff --git a/diploma_thesis/agents/utils/nn/neural_network.py b/diploma_thesis/agents/utils/nn/neural_network.py
@@ -63,11 +63,7 @@ def forward(self, state):
         if isinstance(state, GraphState) and self.graph_encoder is not None:
             data = state.graph.data
 
-            if not self.is_configured:
-                if isinstance(data, HeteroData):
-                    self.graph_encoder = to_hetero(self.graph_encoder, data.metadata(), aggr='sum')
-
-            encoded_graph = self.graph_encoder(data.x_dict, data.edge_index_dict)
+            encoded_graph = self.graph_encoder(data)
 
         hidden = self.merge(encoded_state, encoded_graph)
         output = self.output(hidden)