improve efficiency data module

ThermalSolver/module.py

@@ -111,6 +111,7 @@ class GraphSolver(LightningModule):
             unrolling_steps=self.unrolling_steps,
             batch=batch.batch,
             pos=batch.pos,
+            plot_results=True,
         )
         loss = self._compute_loss(y_pred, y)
         self._log_loss(loss, batch, "test")