add testing

run.py

@@ -1,24 +1,38 @@
+import torch
 from lightning import Trainer
 from ThermalSolver.module import GraphSolver
 from ThermalSolver.data_module import GraphDataModule
 from ThermalSolver.model.local_gno import GatingGNO
-from ThermalSolver.model.basic_gno import GNO
 
 
 def main():
-    trainer = Trainer(max_epochs=100, accelerator="cuda", devices=1)
+    trainer = Trainer(
+        max_epochs=50, accelerator="cuda", devices=1, accumulate_grad_batches=3
+    )
     data_module = GraphDataModule(
         hf_repo="SISSAmathLab/thermal-conduction",
-        split_name="easy",
+        split_name="2000",
         train_size=0.8,
         val_size=0.1,
         test_size=0.1,
-        batch_size=8,
+        batch_size=10,
     )
-    model = GatingGNO(x_ch_node=1, f_ch_node=1, hidden=16, layers=8, edge_ch=3, out_ch=1)
+    data_module.prepare_data()
-    solver = GraphSolver(model)
+    data_module.setup("fit")
-    trainer.fit(solver, datamodule=data_module)
+    model = GatingGNO(
-    print("Done!")
+        x_ch_node=1, f_ch_node=1, hidden=16, layers=2, edge_ch=3, out_ch=1
+    )
+    solver = GraphSolver(model, unrolling_steps=64)
+
+    trainer.fit(
+        solver,
+        train_dataloaders=data_module.train_dataloader(),
+        val_dataloaders=data_module.val_dataloader(),
+    )
+    data_module.setup("test")
+    trainer.test(solver, dataloaders=data_module.test_dataloader())
+
 
 if __name__ == "__main__":
+    torch.set_float32_matmul_precision("medium")
     main()