Add kwargs to MGN forward. (#701)

NVIDIA · Oct 29, 2024 · f7a5773 · f7a5773
1 parent 29774fe
commit f7a5773
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 0 deletions.
diff --git a/examples/cfd/aero_graph_net/models.py b/examples/cfd/aero_graph_net/models.py
@@ -88,6 +88,7 @@ def forward(
         node_features: Tensor,
         edge_features: Tensor,
         graph: Union[DGLGraph, list[DGLGraph], "CuGraphCSC"],
+        **kwargs,
     ) -> Tensor:
         edge_features = self.edge_encoder(edge_features)
         node_features = self.node_encoder(node_features)

diff --git a/modulus/models/meshgraphnet/meshgraphnet.py b/modulus/models/meshgraphnet/meshgraphnet.py
@@ -191,6 +191,7 @@ def forward(
         node_features: Tensor,
         edge_features: Tensor,
         graph: Union[DGLGraph, List[DGLGraph], CuGraphCSC],
+        **kwargs,
     ) -> Tensor:
         edge_features = self.edge_encoder(edge_features)
         node_features = self.node_encoder(node_features)