implement first GNO

ThermalSolver/module.py

@@ -4,7 +4,12 @@ from torch_geometric.data import Batch
 
 
 class GraphSolver(LightningModule):
-    def __init__(self, model: torch.nn.Module, loss: torch.nn.Module = None, unrolling_steps: int = 10):
+    def __init__(
+        self,
+        model: torch.nn.Module,
+        loss: torch.nn.Module = None,
+        unrolling_steps: int = 10,
+    ):
         super().__init__()
         self.model = model
         self.loss = loss if loss is not None else torch.nn.MSELoss()
@@ -18,7 +23,7 @@ class GraphSolver(LightningModule):
         edge_attr: torch.Tensor,
     ):
         return self.model(x, c, edge_index, edge_attr)
-    
+
     def _compute_loss_train(self, x, x_prev, y):
         return self.loss(x, y) + self.loss(x, x_prev)
 
@@ -27,7 +32,7 @@ class GraphSolver(LightningModule):
 
     def _preprocess_batch(self, batch: Batch):
         return batch.x, batch.y, batch.c, batch.edge_index, batch.edge_attr
-    
+
     def _log_loss(self, loss, batch, stage: str):
         self.log(
             f"{stage}_loss",
@@ -41,13 +46,14 @@ class GraphSolver(LightningModule):
 
     def training_step(self, batch: Batch, _):
         x, y, c, edge_index, edge_attr = self._preprocess_batch(batch)
+        loss = 0.0
         for _ in range(self.unrolling_steps):
             x_prev = x.detach()
             x = self(x_prev, c, edge_index=edge_index, edge_attr=edge_attr)
-        loss = self.loss(x, y)
+            loss += self.loss(x, y)
         self._log_loss(loss, batch, "train")
         return loss
-    
+
     def validation_step(self, batch: Batch, _):
         x, y, c, edge_index, edge_attr = self._preprocess_batch(batch)
         for _ in range(self.unrolling_steps):
@@ -70,5 +76,5 @@ class GraphSolver(LightningModule):
         return loss
 
     def configure_optimizers(self):
-        optimizer = torch.optim.Adam(self.parameters(), lr=5e-3)
+        optimizer = torch.optim.Adam(self.parameters(), lr=1e-3)
         return optimizer