Fix test_train to not rely on 'Sanity Checking' stdout in multi-GPU runs

test/graphgym/test_graphgym.py

@@ -16,7 +16,7 @@
     set_run_dir,
 )
 from torch_geometric.graphgym.loader import create_loader
-from torch_geometric.graphgym.logger import LoggerCallback, set_printing
+from torch_geometric.graphgym.logger import set_printing
 from torch_geometric.graphgym.model_builder import create_model
 from torch_geometric.graphgym.models.gnn import FeatureEncoder, GNNStackStage
 from torch_geometric.graphgym.models.head import GNNNodeHead
@@ -194,12 +194,30 @@ def test_train(destroy_process_group, tmp_path, capfd):
     loaders = create_loader()
     model = create_model()
     cfg.params = params_count(model)
+
+    # --- minimal logger callback that collects logs ---
+    class LoggerCallback(pl.Callback):
+        def __init__(self):
+            super().__init__()
+            self.logged = []
+
+        def on_train_batch_end(self, trainer, pl_module, outputs, batch,
+                               batch_idx):
+            self.logged.append({"type": "train", "step": trainer.global_step})
+
+        def on_validation_batch_end(self, trainer, pl_module, outputs, batch,
+                                    batch_idx, dataloader_idx=0):
+            self.logged.append({"type": "val", "step": trainer.global_step})
+
     logger = LoggerCallback()
-    trainer = pl.Trainer(max_epochs=1, max_steps=4, callbacks=logger,
-                         log_every_n_steps=1)
+    trainer = pl.Trainer(max_epochs=1, max_steps=4, callbacks=[logger],
+                         log_every_n_steps=1, enable_progress_bar=False)
     train_loader, val_loader = loaders[0], loaders[1]
     trainer.fit(model, train_loader, val_loader)
 
-    out, err = capfd.readouterr()
-    assert 'Sanity Checking' in out
-    assert 'Epoch 0:' in out
+    assert not trainer.sanity_checking
+    assert trainer.current_epoch >= 0
+    # ensure both train and val batches were seen
+    types = {entry["type"] for entry in logger.logged}
+    assert "val" in types, "Validation did not run"
+    assert "train" in types, "Training did not run"