diff --git a/2026_tdl_challenge/outputs/2026-07-29_18-09-01/results.json b/2026_tdl_challenge/outputs/2026-07-29_18-09-01/results.json
new file mode 100644
index 000000000..a2ba5321f
--- /dev/null
+++ b/2026_tdl_challenge/outputs/2026-07-29_18-09-01/results.json
@@ -0,0 +1,5776 @@
+{
+  "metadata": {
+    "study_id": "2026-07-29_18-09-01",
+    "model_config": "graph/gauge",
+    "generated_at_utc": "2026-07-30T09:36:50.432550+00:00",
+    "n_runs": 72,
+    "train_seeds": [
+      42,
+      43,
+      44
+    ],
+    "heatmap_note": "Cells show mean \u00b1 std over train_seeds (in-distribution test)."
+  },
+  "results": [
+    {
+      "experiment": "community_detection",
+      "wandb_project": "challenge_community_detection",
+      "wandb_run_name": "gauge_hom_0-0.1__deg_1-2.5__gamma_1.5-2__s42",
+      "train_seed": 42,
+      "homophily": "h_lo",
+      "avg_degree": "d_lo",
+      "power_law": "pl_lo",
+      "run_slug": "h_lo__d_lo__pl_lo",
+      "test_loss": 2.6258316040039062,
+      "test_best_rerun_accuracy": 0.2415768951177597,
+      "test_best_rerun_mse": null,
+      "test_triangles_total_structural": null,
+      "test_mse_by_total_triangles": null,
+      "ood_test": {
+        "h_lo__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": 0.23806250095367432,
+          "test_best_rerun_mse": null
+        },
+        "h_lo__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": 0.27901291847229004,
+          "test_best_rerun_mse": null
+        },
+        "h_lo__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": 0.27859270572662354,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": 0.2668271064758301,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": 0.2562457025051117,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": 0.3047597110271454,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": 0.30827412009239197,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": 0.3323401212692261,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": 0.3183589279651642,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": 0.40465277433395386,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": 0.42669418454170227,
+          "test_best_rerun_mse": null
+        }
+      },
+      "output_dir": "/raven/u/klettl/Projects/topobench/logs/train/runs/notebook_gu_grid_2026-07-29_18-09-01__community_detection__00__h_lo__d_lo__pl_lo__s42",
+      "wandb_config": {
+        "AvgTime/train_epoch_mean": 27.257680437781595,
+        "AvgTime/train_epoch_std": 0.09545168304552551,
+        "model/params/total": 682465,
+        "model/params/trainable": 682465,
+        "model/params/non_trainable": 0
+      }
+    },
+    {
+      "experiment": "community_detection",
+      "wandb_project": "challenge_community_detection",
+      "wandb_run_name": "gauge_hom_0-0.1__deg_1-2.5__gamma_1.5-2__s43",
+      "train_seed": 43,
+      "homophily": "h_lo",
+      "avg_degree": "d_lo",
+      "power_law": "pl_lo",
+      "run_slug": "h_lo__d_lo__pl_lo",
+      "test_loss": 2.4898829460144043,
+      "test_best_rerun_accuracy": 0.24493849277496338,
+      "test_best_rerun_mse": null,
+      "test_triangles_total_structural": null,
+      "test_mse_by_total_triangles": null,
+      "ood_test": {
+        "h_lo__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": 0.24146229028701782,
+          "test_best_rerun_mse": null
+        },
+        "h_lo__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": 0.26461151242256165,
+          "test_best_rerun_mse": null
+        },
+        "h_lo__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": 0.262930691242218,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": 0.2525402903556824,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": 0.24585530161857605,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": 0.27439069747924805,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": 0.27546030282974243,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": 0.2955535054206848,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": 0.28760790824890137,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": 0.33596912026405334,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": 0.34987393021583557,
+          "test_best_rerun_mse": null
+        }
+      },
+      "output_dir": "/raven/u/klettl/Projects/topobench/logs/train/runs/notebook_gu_grid_2026-07-29_18-09-01__community_detection__00__h_lo__d_lo__pl_lo__s43",
+      "wandb_config": {
+        "AvgTime/train_epoch_mean": 27.074480921030045,
+        "AvgTime/train_epoch_std": 0.022286526056669996,
+        "model/params/total": 682465,
+        "model/params/trainable": 682465,
+        "model/params/non_trainable": 0
+      }
+    },
+    {
+      "experiment": "community_detection",
+      "wandb_project": "challenge_community_detection",
+      "wandb_run_name": "gauge_hom_0-0.1__deg_1-2.5__gamma_1.5-2__s44",
+      "train_seed": 44,
+      "homophily": "h_lo",
+      "avg_degree": "d_lo",
+      "power_law": "pl_lo",
+      "run_slug": "h_lo__d_lo__pl_lo",
+      "test_loss": 2.609708309173584,
+      "test_best_rerun_accuracy": 0.24165329337120056,
+      "test_best_rerun_mse": null,
+      "test_triangles_total_structural": null,
+      "test_mse_by_total_triangles": null,
+      "ood_test": {
+        "h_lo__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": 0.23726029694080353,
+          "test_best_rerun_mse": null
+        },
+        "h_lo__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": 0.28103750944137573,
+          "test_best_rerun_mse": null
+        },
+        "h_lo__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": 0.28420811891555786,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": 0.27752310037612915,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": 0.2612881064414978,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": 0.3145389258861542,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": 0.32114753127098083,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": 0.3474291265010834,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": 0.34101152420043945,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": 0.4232943654060364,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": 0.4499961733818054,
+          "test_best_rerun_mse": null
+        }
+      },
+      "output_dir": "/raven/u/klettl/Projects/topobench/logs/train/runs/notebook_gu_grid_2026-07-29_18-09-01__community_detection__00__h_lo__d_lo__pl_lo__s44",
+      "wandb_config": {
+        "AvgTime/train_epoch_mean": 27.265852842066025,
+        "AvgTime/train_epoch_std": 1.1835400668804723,
+        "model/params/total": 682465,
+        "model/params/trainable": 682465,
+        "model/params/non_trainable": 0
+      }
+    },
+    {
+      "experiment": "community_detection",
+      "wandb_project": "challenge_community_detection",
+      "wandb_run_name": "gauge_hom_0-0.1__deg_1-2.5__gamma_4-5__s42",
+      "train_seed": 42,
+      "homophily": "h_lo",
+      "avg_degree": "d_lo",
+      "power_law": "pl_hi",
+      "run_slug": "h_lo__d_lo__pl_hi",
+      "test_loss": 2.506941080093384,
+      "test_best_rerun_accuracy": 0.23607610166072845,
+      "test_best_rerun_mse": null,
+      "test_triangles_total_structural": null,
+      "test_mse_by_total_triangles": null,
+      "ood_test": {
+        "h_lo__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": 0.23588509857654572,
+          "test_best_rerun_mse": null
+        },
+        "h_lo__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": 0.24612270295619965,
+          "test_best_rerun_mse": null
+        },
+        "h_lo__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": 0.2444418966770172,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": 0.23768049478530884,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": 0.2323324978351593,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": 0.2432194948196411,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": 0.24562609195709229,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": 0.25036290287971497,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": 0.24745969474315643,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": 0.26464971899986267,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": 0.2708381116390228,
+          "test_best_rerun_mse": null
+        }
+      },
+      "output_dir": "/raven/u/klettl/Projects/topobench/logs/train/runs/notebook_gu_grid_2026-07-29_18-09-01__community_detection__01__h_lo__d_lo__pl_hi__s42",
+      "wandb_config": {
+        "AvgTime/train_epoch_mean": 27.27454023361206,
+        "AvgTime/train_epoch_std": 0.02767995412888662,
+        "model/params/total": 682465,
+        "model/params/trainable": 682465,
+        "model/params/non_trainable": 0
+      }
+    },
+    {
+      "experiment": "community_detection",
+      "wandb_project": "challenge_community_detection",
+      "wandb_run_name": "gauge_hom_0-0.1__deg_1-2.5__gamma_4-5__s43",
+      "train_seed": 43,
+      "homophily": "h_lo",
+      "avg_degree": "d_lo",
+      "power_law": "pl_hi",
+      "run_slug": "h_lo__d_lo__pl_hi",
+      "test_loss": 2.5061724185943604,
+      "test_best_rerun_accuracy": 0.23527389764785767,
+      "test_best_rerun_mse": null,
+      "test_triangles_total_structural": null,
+      "test_mse_by_total_triangles": null,
+      "ood_test": {
+        "h_lo__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": 0.23756588995456696,
+          "test_best_rerun_mse": null
+        },
+        "h_lo__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": 0.24910229444503784,
+          "test_best_rerun_mse": null
+        },
+        "h_lo__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": 0.2514707148075104,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": 0.2415768951177597,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": 0.23752769827842712,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": 0.25403010845184326,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": 0.2550233006477356,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": 0.26827871799468994,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": 0.26583391427993774,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": 0.2932615280151367,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": 0.3034227192401886,
+          "test_best_rerun_mse": null
+        }
+      },
+      "output_dir": "/raven/u/klettl/Projects/topobench/logs/train/runs/notebook_gu_grid_2026-07-29_18-09-01__community_detection__01__h_lo__d_lo__pl_hi__s43",
+      "wandb_config": {
+        "AvgTime/train_epoch_mean": 27.13674637476603,
+        "AvgTime/train_epoch_std": 0.030227690298112682,
+        "model/params/total": 682465,
+        "model/params/trainable": 682465,
+        "model/params/non_trainable": 0
+      }
+    },
+    {
+      "experiment": "community_detection",
+      "wandb_project": "challenge_community_detection",
+      "wandb_run_name": "gauge_hom_0-0.1__deg_1-2.5__gamma_4-5__s44",
+      "train_seed": 44,
+      "homophily": "h_lo",
+      "avg_degree": "d_lo",
+      "power_law": "pl_hi",
+      "run_slug": "h_lo__d_lo__pl_hi",
+      "test_loss": 2.4848856925964355,
+      "test_best_rerun_accuracy": 0.24039269983768463,
+      "test_best_rerun_mse": null,
+      "test_triangles_total_structural": null,
+      "test_mse_by_total_triangles": null,
+      "ood_test": {
+        "h_lo__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": 0.24436549842357635,
+          "test_best_rerun_mse": null
+        },
+        "h_lo__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": 0.24894949793815613,
+          "test_best_rerun_mse": null
+        },
+        "h_lo__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": 0.2513178884983063,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": 0.23672549426555634,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": 0.2342424988746643,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": 0.24276109039783478,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": 0.24558790028095245,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": 0.24830010533332825,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": 0.24776530265808105,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": 0.26201391220092773,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": 0.266750693321228,
+          "test_best_rerun_mse": null
+        }
+      },
+      "output_dir": "/raven/u/klettl/Projects/topobench/logs/train/runs/notebook_gu_grid_2026-07-29_18-09-01__community_detection__01__h_lo__d_lo__pl_hi__s44",
+      "wandb_config": {
+        "AvgTime/train_epoch_mean": 30.498333666059708,
+        "AvgTime/train_epoch_std": 0.12797726465561574,
+        "model/params/total": 682465,
+        "model/params/trainable": 682465,
+        "model/params/non_trainable": 0
+      }
+    },
+    {
+      "experiment": "community_detection",
+      "wandb_project": "challenge_community_detection",
+      "wandb_run_name": "gauge_hom_0-0.1__deg_4-5__gamma_1.5-2__s42",
+      "train_seed": 42,
+      "homophily": "h_lo",
+      "avg_degree": "d_hi",
+      "power_law": "pl_lo",
+      "run_slug": "h_lo__d_hi__pl_lo",
+      "test_loss": 2.4067656993865967,
+      "test_best_rerun_accuracy": 0.29574450850486755,
+      "test_best_rerun_mse": null,
+      "test_triangles_total_structural": null,
+      "test_mse_by_total_triangles": null,
+      "ood_test": {
+        "h_lo__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": 0.23019328713417053,
+          "test_best_rerun_mse": null
+        },
+        "h_lo__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": 0.2202230840921402,
+          "test_best_rerun_mse": null
+        },
+        "h_lo__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": 0.29165711998939514,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": 0.2575063109397888,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": 0.23917029798030853,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": 0.323515921831131,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": 0.32894033193588257,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": 0.3434945344924927,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": 0.32317212224006653,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": 0.45790359377861023,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": 0.4848727881908417,
+          "test_best_rerun_mse": null
+        }
+      },
+      "output_dir": "/raven/u/klettl/Projects/topobench/logs/train/runs/notebook_gu_grid_2026-07-29_18-09-01__community_detection__02__h_lo__d_hi__pl_lo__s42",
+      "wandb_config": {
+        "AvgTime/train_epoch_mean": 27.4669534806852,
+        "AvgTime/train_epoch_std": 0.060453024902884206,
+        "model/params/total": 682465,
+        "model/params/trainable": 682465,
+        "model/params/non_trainable": 0
+      }
+    },
+    {
+      "experiment": "community_detection",
+      "wandb_project": "challenge_community_detection",
+      "wandb_run_name": "gauge_hom_0-0.1__deg_4-5__gamma_1.5-2__s43",
+      "train_seed": 43,
+      "homophily": "h_lo",
+      "avg_degree": "d_hi",
+      "power_law": "pl_lo",
+      "run_slug": "h_lo__d_hi__pl_lo",
+      "test_loss": 2.4102325439453125,
+      "test_best_rerun_accuracy": 0.30800673365592957,
+      "test_best_rerun_mse": null,
+      "test_triangles_total_structural": null,
+      "test_mse_by_total_triangles": null,
+      "ood_test": {
+        "h_lo__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": 0.24066010117530823,
+          "test_best_rerun_mse": null
+        },
+        "h_lo__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": 0.22534188628196716,
+          "test_best_rerun_mse": null
+        },
+        "h_lo__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": 0.29440751671791077,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": 0.27683550119400024,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": 0.25116509199142456,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": 0.33696234226226807,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": 0.3423103392124176,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": 0.3734433352947235,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": 0.35640615224838257,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": 0.4850637912750244,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": 0.5166934132575989,
+          "test_best_rerun_mse": null
+        }
+      },
+      "output_dir": "/raven/u/klettl/Projects/topobench/logs/train/runs/notebook_gu_grid_2026-07-29_18-09-01__community_detection__02__h_lo__d_hi__pl_lo__s43",
+      "wandb_config": {
+        "AvgTime/train_epoch_mean": 27.195537214279174,
+        "AvgTime/train_epoch_std": 0.043776358949474514,
+        "model/params/total": 682465,
+        "model/params/trainable": 682465,
+        "model/params/non_trainable": 0
+      }
+    },
+    {
+      "experiment": "community_detection",
+      "wandb_project": "challenge_community_detection",
+      "wandb_run_name": "gauge_hom_0-0.1__deg_4-5__gamma_1.5-2__s44",
+      "train_seed": 44,
+      "homophily": "h_lo",
+      "avg_degree": "d_hi",
+      "power_law": "pl_lo",
+      "run_slug": "h_lo__d_hi__pl_lo",
+      "test_loss": 2.542980909347534,
+      "test_best_rerun_accuracy": 0.2856215238571167,
+      "test_best_rerun_mse": null,
+      "test_triangles_total_structural": null,
+      "test_mse_by_total_triangles": null,
+      "ood_test": {
+        "h_lo__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": 0.2266024947166443,
+          "test_best_rerun_mse": null
+        },
+        "h_lo__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": 0.2074642777442932,
+          "test_best_rerun_mse": null
+        },
+        "h_lo__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": 0.28978532552719116,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": 0.25548169016838074,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": 0.23088088631629944,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": 0.31652534008026123,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": 0.3240889310836792,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": 0.3435327410697937,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": 0.3297043442726135,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": 0.45412176847457886,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": 0.4957979917526245,
+          "test_best_rerun_mse": null
+        }
+      },
+      "output_dir": "/raven/u/klettl/Projects/topobench/logs/train/runs/notebook_gu_grid_2026-07-29_18-09-01__community_detection__02__h_lo__d_hi__pl_lo__s44",
+      "wandb_config": {
+        "AvgTime/train_epoch_mean": 28.867672937672314,
+        "AvgTime/train_epoch_std": 1.5976765508686417,
+        "model/params/total": 682465,
+        "model/params/trainable": 682465,
+        "model/params/non_trainable": 0
+      }
+    },
+    {
+      "experiment": "community_detection",
+      "wandb_project": "challenge_community_detection",
+      "wandb_run_name": "gauge_hom_0-0.1__deg_4-5__gamma_4-5__s42",
+      "train_seed": 42,
+      "homophily": "h_lo",
+      "avg_degree": "d_hi",
+      "power_law": "pl_hi",
+      "run_slug": "h_lo__d_hi__pl_hi",
+      "test_loss": 2.4422221183776855,
+      "test_best_rerun_accuracy": 0.2859271168708801,
+      "test_best_rerun_mse": null,
+      "test_triangles_total_structural": null,
+      "test_mse_by_total_triangles": null,
+      "ood_test": {
+        "h_lo__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": 0.2315302938222885,
+          "test_best_rerun_mse": null
+        },
+        "h_lo__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": 0.2252654880285263,
+          "test_best_rerun_mse": null
+        },
+        "h_lo__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": 0.2796241044998169,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": 0.2636183202266693,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": 0.24975170195102692,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": 0.3132019340991974,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": 0.33463212847709656,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": 0.3610665500164032,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": 0.36343494057655334,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": 0.47421500086784363,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": 0.5265489816665649,
+          "test_best_rerun_mse": null
+        }
+      },
+      "output_dir": "/raven/u/klettl/Projects/topobench/logs/train/runs/notebook_gu_grid_2026-07-29_18-09-01__community_detection__03__h_lo__d_hi__pl_hi__s42",
+      "wandb_config": {
+        "AvgTime/train_epoch_mean": 29.143832949491646,
+        "AvgTime/train_epoch_std": 1.5830000544339766,
+        "model/params/total": 682465,
+        "model/params/trainable": 682465,
+        "model/params/non_trainable": 0
+      }
+    },
+    {
+      "experiment": "community_detection",
+      "wandb_project": "challenge_community_detection",
+      "wandb_run_name": "gauge_hom_0-0.1__deg_4-5__gamma_4-5__s43",
+      "train_seed": 43,
+      "homophily": "h_lo",
+      "avg_degree": "d_hi",
+      "power_law": "pl_hi",
+      "run_slug": "h_lo__d_hi__pl_hi",
+      "test_loss": 2.3692705631256104,
+      "test_best_rerun_accuracy": 0.29612651467323303,
+      "test_best_rerun_mse": null,
+      "test_triangles_total_structural": null,
+      "test_mse_by_total_triangles": null,
+      "ood_test": {
+        "h_lo__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": 0.24108029901981354,
+          "test_best_rerun_mse": null
+        },
+        "h_lo__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": 0.22984948754310608,
+          "test_best_rerun_mse": null
+        },
+        "h_lo__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": 0.2860035002231598,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": 0.2714492976665497,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": 0.2524639070034027,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": 0.3219115436077118,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": 0.33298954367637634,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": 0.36549773812294006,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": 0.3623271584510803,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": 0.46989840269088745,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": 0.5175719857215881,
+          "test_best_rerun_mse": null
+        }
+      },
+      "output_dir": "/raven/u/klettl/Projects/topobench/logs/train/runs/notebook_gu_grid_2026-07-29_18-09-01__community_detection__03__h_lo__d_hi__pl_hi__s43",
+      "wandb_config": {
+        "AvgTime/train_epoch_mean": 27.339614421129227,
+        "AvgTime/train_epoch_std": 0.14707338040216938,
+        "model/params/total": 682465,
+        "model/params/trainable": 682465,
+        "model/params/non_trainable": 0
+      }
+    },
+    {
+      "experiment": "community_detection",
+      "wandb_project": "challenge_community_detection",
+      "wandb_run_name": "gauge_hom_0-0.1__deg_4-5__gamma_4-5__s44",
+      "train_seed": 44,
+      "homophily": "h_lo",
+      "avg_degree": "d_hi",
+      "power_law": "pl_hi",
+      "run_slug": "h_lo__d_hi__pl_hi",
+      "test_loss": 2.4441046714782715,
+      "test_best_rerun_accuracy": 0.290358304977417,
+      "test_best_rerun_mse": null,
+      "test_triangles_total_structural": null,
+      "test_mse_by_total_triangles": null,
+      "ood_test": {
+        "h_lo__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": 0.22977308928966522,
+          "test_best_rerun_mse": null
+        },
+        "h_lo__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": 0.22400489449501038,
+          "test_best_rerun_mse": null
+        },
+        "h_lo__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": 0.28420811891555786,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": 0.26461151242256165,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": 0.24493849277496338,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": 0.3205745220184326,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": 0.3340209424495697,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": 0.36614716053009033,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": 0.35258615016937256,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": 0.4623347818851471,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": 0.5203605890274048,
+          "test_best_rerun_mse": null
+        }
+      },
+      "output_dir": "/raven/u/klettl/Projects/topobench/logs/train/runs/notebook_gu_grid_2026-07-29_18-09-01__community_detection__03__h_lo__d_hi__pl_hi__s44",
+      "wandb_config": {
+        "AvgTime/train_epoch_mean": 27.385082586606345,
+        "AvgTime/train_epoch_std": 0.08494274428289336,
+        "model/params/total": 682465,
+        "model/params/trainable": 682465,
+        "model/params/non_trainable": 0
+      }
+    },
+    {
+      "experiment": "community_detection",
+      "wandb_project": "challenge_community_detection",
+      "wandb_run_name": "gauge_hom_0.4-0.6__deg_1-2.5__gamma_1.5-2__s42",
+      "train_seed": 42,
+      "homophily": "h_mid",
+      "avg_degree": "d_lo",
+      "power_law": "pl_lo",
+      "run_slug": "h_mid__d_lo__pl_lo",
+      "test_loss": 2.346010446548462,
+      "test_best_rerun_accuracy": 0.3542287349700928,
+      "test_best_rerun_mse": null,
+      "test_triangles_total_structural": null,
+      "test_mse_by_total_triangles": null,
+      "ood_test": {
+        "h_lo__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": 0.21346168220043182,
+          "test_best_rerun_mse": null
+        },
+        "h_lo__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": 0.20543968677520752,
+          "test_best_rerun_mse": null
+        },
+        "h_lo__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": 0.23240889608860016,
+          "test_best_rerun_mse": null
+        },
+        "h_lo__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": 0.24467109143733978,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": 0.3247383236885071,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": 0.37959355115890503,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": 0.40148216485977173,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": 0.5129116177558899,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": 0.5069141983985901,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": 0.6067308187484741,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": 0.6450072526931763,
+          "test_best_rerun_mse": null
+        }
+      },
+      "output_dir": "/raven/u/klettl/Projects/topobench/logs/train/runs/notebook_gu_grid_2026-07-29_18-09-01__community_detection__04__h_mid__d_lo__pl_lo__s42",
+      "wandb_config": {
+        "AvgTime/train_epoch_mean": 28.926682436907733,
+        "AvgTime/train_epoch_std": 1.5761604427914102,
+        "model/params/total": 682465,
+        "model/params/trainable": 682465,
+        "model/params/non_trainable": 0
+      }
+    },
+    {
+      "experiment": "community_detection",
+      "wandb_project": "challenge_community_detection",
+      "wandb_run_name": "gauge_hom_0.4-0.6__deg_1-2.5__gamma_1.5-2__s43",
+      "train_seed": 43,
+      "homophily": "h_mid",
+      "avg_degree": "d_lo",
+      "power_law": "pl_lo",
+      "run_slug": "h_mid__d_lo__pl_lo",
+      "test_loss": 2.3754777908325195,
+      "test_best_rerun_accuracy": 0.3494919538497925,
+      "test_best_rerun_mse": null,
+      "test_triangles_total_structural": null,
+      "test_mse_by_total_triangles": null,
+      "ood_test": {
+        "h_lo__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": 0.21812207996845245,
+          "test_best_rerun_mse": null
+        },
+        "h_lo__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": 0.20647108554840088,
+          "test_best_rerun_mse": null
+        },
+        "h_lo__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": 0.22801588475704193,
+          "test_best_rerun_mse": null
+        },
+        "h_lo__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": 0.23718389868736267,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": 0.32661011815071106,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": 0.3670639395713806,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": 0.38643136620521545,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": 0.5204752087593079,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": 0.5118420124053955,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": 0.6085262298583984,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": 0.6430972814559937,
+          "test_best_rerun_mse": null
+        }
+      },
+      "output_dir": "/raven/u/klettl/Projects/topobench/logs/train/runs/notebook_gu_grid_2026-07-29_18-09-01__community_detection__04__h_mid__d_lo__pl_lo__s43",
+      "wandb_config": {
+        "AvgTime/train_epoch_mean": 26.917261355263847,
+        "AvgTime/train_epoch_std": 0.1459446104336187,
+        "model/params/total": 682465,
+        "model/params/trainable": 682465,
+        "model/params/non_trainable": 0
+      }
+    },
+    {
+      "experiment": "community_detection",
+      "wandb_project": "challenge_community_detection",
+      "wandb_run_name": "gauge_hom_0.4-0.6__deg_1-2.5__gamma_1.5-2__s44",
+      "train_seed": 44,
+      "homophily": "h_mid",
+      "avg_degree": "d_lo",
+      "power_law": "pl_lo",
+      "run_slug": "h_mid__d_lo__pl_lo",
+      "test_loss": 2.2358174324035645,
+      "test_best_rerun_accuracy": 0.35323554277420044,
+      "test_best_rerun_mse": null,
+      "test_triangles_total_structural": null,
+      "test_mse_by_total_triangles": null,
+      "ood_test": {
+        "h_lo__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": 0.21774008870124817,
+          "test_best_rerun_mse": null
+        },
+        "h_lo__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": 0.2052868753671646,
+          "test_best_rerun_mse": null
+        },
+        "h_lo__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": 0.22297349572181702,
+          "test_best_rerun_mse": null
+        },
+        "h_lo__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": 0.2296202927827835,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": 0.3251585364341736,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": 0.3689357340335846,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": 0.3888379633426666,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": 0.5062265992164612,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": 0.4982045888900757,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": 0.5943158268928528,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": 0.6222018599510193,
+          "test_best_rerun_mse": null
+        }
+      },
+      "output_dir": "/raven/u/klettl/Projects/topobench/logs/train/runs/notebook_gu_grid_2026-07-29_18-09-01__community_detection__04__h_mid__d_lo__pl_lo__s44",
+      "wandb_config": {
+        "AvgTime/train_epoch_mean": 26.941870898008347,
+        "AvgTime/train_epoch_std": 0.11651695122224251,
+        "model/params/total": 682465,
+        "model/params/trainable": 682465,
+        "model/params/non_trainable": 0
+      }
+    },
+    {
+      "experiment": "community_detection",
+      "wandb_project": "challenge_community_detection",
+      "wandb_run_name": "gauge_hom_0.4-0.6__deg_1-2.5__gamma_4-5__s42",
+      "train_seed": 42,
+      "homophily": "h_mid",
+      "avg_degree": "d_lo",
+      "power_law": "pl_hi",
+      "run_slug": "h_mid__d_lo__pl_hi",
+      "test_loss": 2.4543328285217285,
+      "test_best_rerun_accuracy": 0.3283291459083557,
+      "test_best_rerun_mse": null,
+      "test_triangles_total_structural": null,
+      "test_mse_by_total_triangles": null,
+      "ood_test": {
+        "h_lo__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": 0.20899228751659393,
+          "test_best_rerun_mse": null
+        },
+        "h_lo__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": 0.20177248120307922,
+          "test_best_rerun_mse": null
+        },
+        "h_lo__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": 0.23439529538154602,
+          "test_best_rerun_mse": null
+        },
+        "h_lo__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": 0.2499808967113495,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": 0.34303614497184753,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": 0.370005339384079,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": 0.4113377630710602,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": 0.503055989742279,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": 0.5041255950927734,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": 0.5940484404563904,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": 0.6451600790023804,
+          "test_best_rerun_mse": null
+        }
+      },
+      "output_dir": "/raven/u/klettl/Projects/topobench/logs/train/runs/notebook_gu_grid_2026-07-29_18-09-01__community_detection__05__h_mid__d_lo__pl_hi__s42",
+      "wandb_config": {
+        "AvgTime/train_epoch_mean": 26.750238127178616,
+        "AvgTime/train_epoch_std": 0.059019451113226276,
+        "model/params/total": 682465,
+        "model/params/trainable": 682465,
+        "model/params/non_trainable": 0
+      }
+    },
+    {
+      "experiment": "community_detection",
+      "wandb_project": "challenge_community_detection",
+      "wandb_run_name": "gauge_hom_0.4-0.6__deg_1-2.5__gamma_4-5__s43",
+      "train_seed": 43,
+      "homophily": "h_mid",
+      "avg_degree": "d_lo",
+      "power_law": "pl_hi",
+      "run_slug": "h_mid__d_lo__pl_hi",
+      "test_loss": 2.389476776123047,
+      "test_best_rerun_accuracy": 0.33054473996162415,
+      "test_best_rerun_mse": null,
+      "test_triangles_total_structural": null,
+      "test_mse_by_total_triangles": null,
+      "ood_test": {
+        "h_lo__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": 0.21667048335075378,
+          "test_best_rerun_mse": null
+        },
+        "h_lo__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": 0.2037970870733261,
+          "test_best_rerun_mse": null
+        },
+        "h_lo__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": 0.2252654880285263,
+          "test_best_rerun_mse": null
+        },
+        "h_lo__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": 0.24092750251293182,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": 0.3419283330440521,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": 0.3736725449562073,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": 0.40980976819992065,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": 0.5138283967971802,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": 0.5187562108039856,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": 0.6075712442398071,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": 0.659408688545227,
+          "test_best_rerun_mse": null
+        }
+      },
+      "output_dir": "/raven/u/klettl/Projects/topobench/logs/train/runs/notebook_gu_grid_2026-07-29_18-09-01__community_detection__05__h_mid__d_lo__pl_hi__s43",
+      "wandb_config": {
+        "AvgTime/train_epoch_mean": 27.12549352645874,
+        "AvgTime/train_epoch_std": 0.053675788459561204,
+        "model/params/total": 682465,
+        "model/params/trainable": 682465,
+        "model/params/non_trainable": 0
+      }
+    },
+    {
+      "experiment": "community_detection",
+      "wandb_project": "challenge_community_detection",
+      "wandb_run_name": "gauge_hom_0.4-0.6__deg_1-2.5__gamma_4-5__s44",
+      "train_seed": 44,
+      "homophily": "h_mid",
+      "avg_degree": "d_lo",
+      "power_law": "pl_hi",
+      "run_slug": "h_mid__d_lo__pl_hi",
+      "test_loss": 2.371619701385498,
+      "test_best_rerun_accuracy": 0.33096492290496826,
+      "test_best_rerun_mse": null,
+      "test_triangles_total_structural": null,
+      "test_mse_by_total_triangles": null,
+      "ood_test": {
+        "h_lo__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": 0.21258307993412018,
+          "test_best_rerun_mse": null
+        },
+        "h_lo__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": 0.20692948997020721,
+          "test_best_rerun_mse": null
+        },
+        "h_lo__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": 0.22133088111877441,
+          "test_best_rerun_mse": null
+        },
+        "h_lo__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": 0.2364199012517929,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": 0.34616854786872864,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": 0.3672931492328644,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": 0.39854076504707336,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": 0.5132172107696533,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": 0.5124149918556213,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": 0.598365068435669,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": 0.6463442444801331,
+          "test_best_rerun_mse": null
+        }
+      },
+      "output_dir": "/raven/u/klettl/Projects/topobench/logs/train/runs/notebook_gu_grid_2026-07-29_18-09-01__community_detection__05__h_mid__d_lo__pl_hi__s44",
+      "wandb_config": {
+        "AvgTime/train_epoch_mean": 29.42714535196622,
+        "AvgTime/train_epoch_std": 1.6148785477369063,
+        "model/params/total": 682465,
+        "model/params/trainable": 682465,
+        "model/params/non_trainable": 0
+      }
+    },
+    {
+      "experiment": "community_detection",
+      "wandb_project": "challenge_community_detection",
+      "wandb_run_name": "gauge_hom_0.4-0.6__deg_4-5__gamma_1.5-2__s42",
+      "train_seed": 42,
+      "homophily": "h_mid",
+      "avg_degree": "d_hi",
+      "power_law": "pl_lo",
+      "run_slug": "h_mid__d_hi__pl_lo",
+      "test_loss": 2.286994457244873,
+      "test_best_rerun_accuracy": 0.3967835605144501,
+      "test_best_rerun_mse": null,
+      "test_triangles_total_structural": null,
+      "test_mse_by_total_triangles": null,
+      "ood_test": {
+        "h_lo__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": 0.19531667232513428,
+          "test_best_rerun_mse": null
+        },
+        "h_lo__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": 0.18630146980285645,
+          "test_best_rerun_mse": null
+        },
+        "h_lo__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": 0.2396668940782547,
+          "test_best_rerun_mse": null
+        },
+        "h_lo__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": 0.2447092980146408,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": 0.31507372856140137,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": 0.28756970167160034,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": 0.41156697273254395,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": 0.4661547839641571,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": 0.4546183943748474,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": 0.5993200540542603,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": 0.6339674592018127,
+          "test_best_rerun_mse": null
+        }
+      },
+      "output_dir": "/raven/u/klettl/Projects/topobench/logs/train/runs/notebook_gu_grid_2026-07-29_18-09-01__community_detection__06__h_mid__d_hi__pl_lo__s42",
+      "wandb_config": {
+        "AvgTime/train_epoch_mean": 29.394388556480408,
+        "AvgTime/train_epoch_std": 1.7454243628371895,
+        "model/params/total": 682465,
+        "model/params/trainable": 682465,
+        "model/params/non_trainable": 0
+      }
+    },
+    {
+      "experiment": "community_detection",
+      "wandb_project": "challenge_community_detection",
+      "wandb_run_name": "gauge_hom_0.4-0.6__deg_4-5__gamma_1.5-2__s43",
+      "train_seed": 43,
+      "homophily": "h_mid",
+      "avg_degree": "d_hi",
+      "power_law": "pl_lo",
+      "run_slug": "h_mid__d_hi__pl_lo",
+      "test_loss": 2.214320659637451,
+      "test_best_rerun_accuracy": 0.396707147359848,
+      "test_best_rerun_mse": null,
+      "test_triangles_total_structural": null,
+      "test_mse_by_total_triangles": null,
+      "ood_test": {
+        "h_lo__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": 0.19252808392047882,
+          "test_best_rerun_mse": null
+        },
+        "h_lo__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": 0.18484987318515778,
+          "test_best_rerun_mse": null
+        },
+        "h_lo__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": 0.24367789924144745,
+          "test_best_rerun_mse": null
+        },
+        "h_lo__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": 0.24352510273456573,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": 0.3261135220527649,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": 0.2939109206199646,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": 0.4170677661895752,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": 0.4862861931324005,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": 0.46691879630088806,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": 0.6063488721847534,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": 0.6405760645866394,
+          "test_best_rerun_mse": null
+        }
+      },
+      "output_dir": "/raven/u/klettl/Projects/topobench/logs/train/runs/notebook_gu_grid_2026-07-29_18-09-01__community_detection__06__h_mid__d_hi__pl_lo__s43",
+      "wandb_config": {
+        "AvgTime/train_epoch_mean": 27.477708803979976,
+        "AvgTime/train_epoch_std": 0.338585928561687,
+        "model/params/total": 682465,
+        "model/params/trainable": 682465,
+        "model/params/non_trainable": 0
+      }
+    },
+    {
+      "experiment": "community_detection",
+      "wandb_project": "challenge_community_detection",
+      "wandb_run_name": "gauge_hom_0.4-0.6__deg_4-5__gamma_1.5-2__s44",
+      "train_seed": 44,
+      "homophily": "h_mid",
+      "avg_degree": "d_hi",
+      "power_law": "pl_lo",
+      "run_slug": "h_mid__d_hi__pl_lo",
+      "test_loss": 2.1461167335510254,
+      "test_best_rerun_accuracy": 0.39731836318969727,
+      "test_best_rerun_mse": null,
+      "test_triangles_total_structural": null,
+      "test_mse_by_total_triangles": null,
+      "ood_test": {
+        "h_lo__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": 0.18931928277015686,
+          "test_best_rerun_mse": null
+        },
+        "h_lo__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": 0.18358927965164185,
+          "test_best_rerun_mse": null
+        },
+        "h_lo__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": 0.24127130210399628,
+          "test_best_rerun_mse": null
+        },
+        "h_lo__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": 0.24295209348201752,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": 0.3222171366214752,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": 0.29219192266464233,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": 0.4129803776741028,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": 0.4784933924674988,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": 0.45927879214286804,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": 0.600542426109314,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": 0.634884238243103,
+          "test_best_rerun_mse": null
+        }
+      },
+      "output_dir": "/raven/u/klettl/Projects/topobench/logs/train/runs/notebook_gu_grid_2026-07-29_18-09-01__community_detection__06__h_mid__d_hi__pl_lo__s44",
+      "wandb_config": {
+        "AvgTime/train_epoch_mean": 29.48760253853268,
+        "AvgTime/train_epoch_std": 1.6977944497328943,
+        "model/params/total": 682465,
+        "model/params/trainable": 682465,
+        "model/params/non_trainable": 0
+      }
+    },
+    {
+      "experiment": "community_detection",
+      "wandb_project": "challenge_community_detection",
+      "wandb_run_name": "gauge_hom_0.4-0.6__deg_4-5__gamma_4-5__s42",
+      "train_seed": 42,
+      "homophily": "h_mid",
+      "avg_degree": "d_hi",
+      "power_law": "pl_hi",
+      "run_slug": "h_mid__d_hi__pl_hi",
+      "test_loss": 2.1196160316467285,
+      "test_best_rerun_accuracy": 0.4402933716773987,
+      "test_best_rerun_mse": null,
+      "test_triangles_total_structural": null,
+      "test_mse_by_total_triangles": null,
+      "ood_test": {
+        "h_lo__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": 0.19172587990760803,
+          "test_best_rerun_mse": null
+        },
+        "h_lo__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": 0.18737107515335083,
+          "test_best_rerun_mse": null
+        },
+        "h_lo__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": 0.22121629118919373,
+          "test_best_rerun_mse": null
+        },
+        "h_lo__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": 0.23099549114704132,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": 0.3398655354976654,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": 0.31816792488098145,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": 0.3958285450935364,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": 0.5098174214363098,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": 0.511345386505127,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": 0.6192222237586975,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": 0.6716326475143433,
+          "test_best_rerun_mse": null
+        }
+      },
+      "output_dir": "/raven/u/klettl/Projects/topobench/logs/train/runs/notebook_gu_grid_2026-07-29_18-09-01__community_detection__07__h_mid__d_hi__pl_hi__s42",
+      "wandb_config": {
+        "AvgTime/train_epoch_mean": 27.28913544654846,
+        "AvgTime/train_epoch_std": 0.17596039311220518,
+        "model/params/total": 682465,
+        "model/params/trainable": 682465,
+        "model/params/non_trainable": 0
+      }
+    },
+    {
+      "experiment": "community_detection",
+      "wandb_project": "challenge_community_detection",
+      "wandb_run_name": "gauge_hom_0.4-0.6__deg_4-5__gamma_4-5__s43",
+      "train_seed": 43,
+      "homophily": "h_mid",
+      "avg_degree": "d_hi",
+      "power_law": "pl_hi",
+      "run_slug": "h_mid__d_hi__pl_hi",
+      "test_loss": 2.103053569793701,
+      "test_best_rerun_accuracy": 0.44040796160697937,
+      "test_best_rerun_mse": null,
+      "test_triangles_total_structural": null,
+      "test_mse_by_total_triangles": null,
+      "ood_test": {
+        "h_lo__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": 0.19894568622112274,
+          "test_best_rerun_mse": null
+        },
+        "h_lo__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": 0.1924898773431778,
+          "test_best_rerun_mse": null
+        },
+        "h_lo__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": 0.22908549010753632,
+          "test_best_rerun_mse": null
+        },
+        "h_lo__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": 0.24123309552669525,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": 0.33803194761276245,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": 0.31816792488098145,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": 0.3969745635986328,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": 0.511498212814331,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": 0.5065321922302246,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": 0.6160516738891602,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": 0.6691114902496338,
+          "test_best_rerun_mse": null
+        }
+      },
+      "output_dir": "/raven/u/klettl/Projects/topobench/logs/train/runs/notebook_gu_grid_2026-07-29_18-09-01__community_detection__07__h_mid__d_hi__pl_hi__s43",
+      "wandb_config": {
+        "AvgTime/train_epoch_mean": 29.16252639180138,
+        "AvgTime/train_epoch_std": 1.740703492372861,
+        "model/params/total": 682465,
+        "model/params/trainable": 682465,
+        "model/params/non_trainable": 0
+      }
+    },
+    {
+      "experiment": "community_detection",
+      "wandb_project": "challenge_community_detection",
+      "wandb_run_name": "gauge_hom_0.4-0.6__deg_4-5__gamma_4-5__s44",
+      "train_seed": 44,
+      "homophily": "h_mid",
+      "avg_degree": "d_hi",
+      "power_law": "pl_hi",
+      "run_slug": "h_mid__d_hi__pl_hi",
+      "test_loss": 2.0471880435943604,
+      "test_best_rerun_accuracy": 0.4427763819694519,
+      "test_best_rerun_mse": null,
+      "test_triangles_total_structural": null,
+      "test_mse_by_total_triangles": null,
+      "ood_test": {
+        "h_lo__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": 0.19673007726669312,
+          "test_best_rerun_mse": null
+        },
+        "h_lo__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": 0.19077087938785553,
+          "test_best_rerun_mse": null
+        },
+        "h_lo__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": 0.22862708568572998,
+          "test_best_rerun_mse": null
+        },
+        "h_lo__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": 0.23806250095367432,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": 0.3373061418533325,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": 0.31656351685523987,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": 0.3957521617412567,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": 0.5095118284225464,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": 0.5060356259346008,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": 0.6134922504425049,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": 0.6723966598510742,
+          "test_best_rerun_mse": null
+        }
+      },
+      "output_dir": "/raven/u/klettl/Projects/topobench/logs/train/runs/notebook_gu_grid_2026-07-29_18-09-01__community_detection__07__h_mid__d_hi__pl_hi__s44",
+      "wandb_config": {
+        "AvgTime/train_epoch_mean": 27.293564376376924,
+        "AvgTime/train_epoch_std": 0.04926143847291967,
+        "model/params/total": 682465,
+        "model/params/trainable": 682465,
+        "model/params/non_trainable": 0
+      }
+    },
+    {
+      "experiment": "community_detection",
+      "wandb_project": "challenge_community_detection",
+      "wandb_run_name": "gauge_hom_0.9-1__deg_1-2.5__gamma_1.5-2__s42",
+      "train_seed": 42,
+      "homophily": "h_hi",
+      "avg_degree": "d_lo",
+      "power_law": "pl_lo",
+      "run_slug": "h_hi__d_lo__pl_lo",
+      "test_loss": 1.714068055152893,
+      "test_best_rerun_accuracy": 0.565589427947998,
+      "test_best_rerun_mse": null,
+      "test_triangles_total_structural": null,
+      "test_mse_by_total_triangles": null,
+      "ood_test": {
+        "h_lo__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": 0.19432348012924194,
+          "test_best_rerun_mse": null
+        },
+        "h_lo__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": 0.18431507050991058,
+          "test_best_rerun_mse": null
+        },
+        "h_lo__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": 0.20169608294963837,
+          "test_best_rerun_mse": null
+        },
+        "h_lo__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": 0.20788449048995972,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": 0.3599587380886078,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": 0.3344411253929138,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": 0.3530445396900177,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": 0.3874245584011078,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": 0.5600504279136658,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": 0.6445106863975525,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": 0.6854228973388672,
+          "test_best_rerun_mse": null
+        }
+      },
+      "output_dir": "/raven/u/klettl/Projects/topobench/logs/train/runs/notebook_gu_grid_2026-07-29_18-09-01__community_detection__08__h_hi__d_lo__pl_lo__s42",
+      "wandb_config": {
+        "AvgTime/train_epoch_mean": 27.144500330090523,
+        "AvgTime/train_epoch_std": 0.9532709104724897,
+        "model/params/total": 682465,
+        "model/params/trainable": 682465,
+        "model/params/non_trainable": 0
+      }
+    },
+    {
+      "experiment": "community_detection",
+      "wandb_project": "challenge_community_detection",
+      "wandb_run_name": "gauge_hom_0.9-1__deg_1-2.5__gamma_1.5-2__s43",
+      "train_seed": 43,
+      "homophily": "h_hi",
+      "avg_degree": "d_lo",
+      "power_law": "pl_lo",
+      "run_slug": "h_hi__d_lo__pl_lo",
+      "test_loss": 1.8442329168319702,
+      "test_best_rerun_accuracy": 0.5592864155769348,
+      "test_best_rerun_mse": null,
+      "test_triangles_total_structural": null,
+      "test_mse_by_total_triangles": null,
+      "ood_test": {
+        "h_lo__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": 0.1888226717710495,
+          "test_best_rerun_mse": null
+        },
+        "h_lo__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": 0.18171747028827667,
+          "test_best_rerun_mse": null
+        },
+        "h_lo__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": 0.16792726516723633,
+          "test_best_rerun_mse": null
+        },
+        "h_lo__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": 0.16655206680297852,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": 0.3499503433704376,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": 0.3299335241317749,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": 0.3225991427898407,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": 0.33436474204063416,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": 0.5606616139411926,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": 0.6333562731742859,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": 0.6655206680297852,
+          "test_best_rerun_mse": null
+        }
+      },
+      "output_dir": "/raven/u/klettl/Projects/topobench/logs/train/runs/notebook_gu_grid_2026-07-29_18-09-01__community_detection__08__h_hi__d_lo__pl_lo__s43",
+      "wandb_config": {
+        "AvgTime/train_epoch_mean": 29.72986540373634,
+        "AvgTime/train_epoch_std": 1.371840279181708,
+        "model/params/total": 682465,
+        "model/params/trainable": 682465,
+        "model/params/non_trainable": 0
+      }
+    },
+    {
+      "experiment": "community_detection",
+      "wandb_project": "challenge_community_detection",
+      "wandb_run_name": "gauge_hom_0.9-1__deg_1-2.5__gamma_1.5-2__s44",
+      "train_seed": 44,
+      "homophily": "h_hi",
+      "avg_degree": "d_lo",
+      "power_law": "pl_lo",
+      "run_slug": "h_hi__d_lo__pl_lo",
+      "test_loss": 1.7324192523956299,
+      "test_best_rerun_accuracy": 0.5631064176559448,
+      "test_best_rerun_mse": null,
+      "test_triangles_total_structural": null,
+      "test_mse_by_total_triangles": null,
+      "ood_test": {
+        "h_lo__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": 0.19600427150726318,
+          "test_best_rerun_mse": null
+        },
+        "h_lo__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": 0.1863778680562973,
+          "test_best_rerun_mse": null
+        },
+        "h_lo__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": 0.1979142725467682,
+          "test_best_rerun_mse": null
+        },
+        "h_lo__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": 0.2066238820552826,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": 0.3560241460800171,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": 0.33868134021759033,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": 0.34807854890823364,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": 0.3843303620815277,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": 0.5662388205528259,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": 0.643250048160553,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": 0.678355872631073,
+          "test_best_rerun_mse": null
+        }
+      },
+      "output_dir": "/raven/u/klettl/Projects/topobench/logs/train/runs/notebook_gu_grid_2026-07-29_18-09-01__community_detection__08__h_hi__d_lo__pl_lo__s44",
+      "wandb_config": {
+        "AvgTime/train_epoch_mean": 27.175156247231268,
+        "AvgTime/train_epoch_std": 0.8513084127345705,
+        "model/params/total": 682465,
+        "model/params/trainable": 682465,
+        "model/params/non_trainable": 0
+      }
+    },
+    {
+      "experiment": "community_detection",
+      "wandb_project": "challenge_community_detection",
+      "wandb_run_name": "gauge_hom_0.9-1__deg_1-2.5__gamma_4-5__s42",
+      "train_seed": 42,
+      "homophily": "h_hi",
+      "avg_degree": "d_lo",
+      "power_law": "pl_hi",
+      "run_slug": "h_hi__d_lo__pl_hi",
+      "test_loss": 1.7422338724136353,
+      "test_best_rerun_accuracy": 0.5794942378997803,
+      "test_best_rerun_mse": null,
+      "test_triangles_total_structural": null,
+      "test_mse_by_total_triangles": null,
+      "ood_test": {
+        "h_lo__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": 0.19218426942825317,
+          "test_best_rerun_mse": null
+        },
+        "h_lo__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": 0.18206126987934113,
+          "test_best_rerun_mse": null
+        },
+        "h_lo__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": 0.17839406430721283,
+          "test_best_rerun_mse": null
+        },
+        "h_lo__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": 0.19283367693424225,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": 0.36274734139442444,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": 0.3447551429271698,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": 0.34345632791519165,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": 0.3858201503753662,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": 0.5728856325149536,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": 0.650928258895874,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": 0.6945144534111023,
+          "test_best_rerun_mse": null
+        }
+      },
+      "output_dir": "/raven/u/klettl/Projects/topobench/logs/train/runs/notebook_gu_grid_2026-07-29_18-09-01__community_detection__09__h_hi__d_lo__pl_hi__s42",
+      "wandb_config": {
+        "AvgTime/train_epoch_mean": 26.88993050654729,
+        "AvgTime/train_epoch_std": 0.0941506305230491,
+        "model/params/total": 682465,
+        "model/params/trainable": 682465,
+        "model/params/non_trainable": 0
+      }
+    },
+    {
+      "experiment": "community_detection",
+      "wandb_project": "challenge_community_detection",
+      "wandb_run_name": "gauge_hom_0.9-1__deg_1-2.5__gamma_4-5__s43",
+      "train_seed": 43,
+      "homophily": "h_hi",
+      "avg_degree": "d_lo",
+      "power_law": "pl_hi",
+      "run_slug": "h_hi__d_lo__pl_hi",
+      "test_loss": 1.6678621768951416,
+      "test_best_rerun_accuracy": 0.5746810436248779,
+      "test_best_rerun_mse": null,
+      "test_triangles_total_structural": null,
+      "test_mse_by_total_triangles": null,
+      "ood_test": {
+        "h_lo__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": 0.18798227608203888,
+          "test_best_rerun_mse": null
+        },
+        "h_lo__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": 0.17751547694206238,
+          "test_best_rerun_mse": null
+        },
+        "h_lo__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": 0.15421345829963684,
+          "test_best_rerun_mse": null
+        },
+        "h_lo__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": 0.15337306261062622,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": 0.3537321388721466,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": 0.3405149281024933,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": 0.3240507245063782,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": 0.3492627441883087,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": 0.5658186078071594,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": 0.6415692567825317,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": 0.6880586743354797,
+          "test_best_rerun_mse": null
+        }
+      },
+      "output_dir": "/raven/u/klettl/Projects/topobench/logs/train/runs/notebook_gu_grid_2026-07-29_18-09-01__community_detection__09__h_hi__d_lo__pl_hi__s43",
+      "wandb_config": {
+        "AvgTime/train_epoch_mean": 30.386446579642918,
+        "AvgTime/train_epoch_std": 0.09022003076221828,
+        "model/params/total": 682465,
+        "model/params/trainable": 682465,
+        "model/params/non_trainable": 0
+      }
+    },
+    {
+      "experiment": "community_detection",
+      "wandb_project": "challenge_community_detection",
+      "wandb_run_name": "gauge_hom_0.9-1__deg_1-2.5__gamma_4-5__s44",
+      "train_seed": 44,
+      "homophily": "h_hi",
+      "avg_degree": "d_lo",
+      "power_law": "pl_hi",
+      "run_slug": "h_hi__d_lo__pl_hi",
+      "test_loss": 1.6377265453338623,
+      "test_best_rerun_accuracy": 0.5784246325492859,
+      "test_best_rerun_mse": null,
+      "test_triangles_total_structural": null,
+      "test_mse_by_total_triangles": null,
+      "ood_test": {
+        "h_lo__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": 0.18928107619285583,
+          "test_best_rerun_mse": null
+        },
+        "h_lo__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": 0.17954006791114807,
+          "test_best_rerun_mse": null
+        },
+        "h_lo__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": 0.1914202719926834,
+          "test_best_rerun_mse": null
+        },
+        "h_lo__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": 0.1979524791240692,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": 0.3584689497947693,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": 0.3397509455680847,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": 0.3428833484649658,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": 0.38914355635643005,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": 0.5652456283569336,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": 0.6480250358581543,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": 0.6953930854797363,
+          "test_best_rerun_mse": null
+        }
+      },
+      "output_dir": "/raven/u/klettl/Projects/topobench/logs/train/runs/notebook_gu_grid_2026-07-29_18-09-01__community_detection__09__h_hi__d_lo__pl_hi__s44",
+      "wandb_config": {
+        "AvgTime/train_epoch_mean": 29.765423002243043,
+        "AvgTime/train_epoch_std": 1.4600357195344895,
+        "model/params/total": 682465,
+        "model/params/trainable": 682465,
+        "model/params/non_trainable": 0
+      }
+    },
+    {
+      "experiment": "community_detection",
+      "wandb_project": "challenge_community_detection",
+      "wandb_run_name": "gauge_hom_0.9-1__deg_4-5__gamma_1.5-2__s42",
+      "train_seed": 42,
+      "homophily": "h_hi",
+      "avg_degree": "d_hi",
+      "power_law": "pl_lo",
+      "run_slug": "h_hi__d_hi__pl_lo",
+      "test_loss": 1.4983386993408203,
+      "test_best_rerun_accuracy": 0.6427916288375854,
+      "test_best_rerun_mse": null,
+      "test_triangles_total_structural": null,
+      "test_mse_by_total_triangles": null,
+      "ood_test": {
+        "h_lo__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": 0.17572006583213806,
+          "test_best_rerun_mse": null
+        },
+        "h_lo__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": 0.1644892692565918,
+          "test_best_rerun_mse": null
+        },
+        "h_lo__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": 0.20020627975463867,
+          "test_best_rerun_mse": null
+        },
+        "h_lo__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": 0.2013522833585739,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": 0.3374207317829132,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": 0.3126671314239502,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": 0.3707311451435089,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": 0.40232256054878235,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": 0.5407212376594543,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": 0.5318205952644348,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": 0.6839712858200073,
+          "test_best_rerun_mse": null
+        }
+      },
+      "output_dir": "/raven/u/klettl/Projects/topobench/logs/train/runs/notebook_gu_grid_2026-07-29_18-09-01__community_detection__10__h_hi__d_hi__pl_lo__s42",
+      "wandb_config": {
+        "AvgTime/train_epoch_mean": 27.76003390153249,
+        "AvgTime/train_epoch_std": 1.191235574146216,
+        "model/params/total": 682465,
+        "model/params/trainable": 682465,
+        "model/params/non_trainable": 0
+      }
+    },
+    {
+      "experiment": "community_detection",
+      "wandb_project": "challenge_community_detection",
+      "wandb_run_name": "gauge_hom_0.9-1__deg_4-5__gamma_1.5-2__s43",
+      "train_seed": 43,
+      "homophily": "h_hi",
+      "avg_degree": "d_hi",
+      "power_law": "pl_lo",
+      "run_slug": "h_hi__d_hi__pl_lo",
+      "test_loss": 1.4317357540130615,
+      "test_best_rerun_accuracy": 0.6464970707893372,
+      "test_best_rerun_mse": null,
+      "test_triangles_total_structural": null,
+      "test_mse_by_total_triangles": null,
+      "ood_test": {
+        "h_lo__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": 0.1833982765674591,
+          "test_best_rerun_mse": null
+        },
+        "h_lo__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": 0.17212927341461182,
+          "test_best_rerun_mse": null
+        },
+        "h_lo__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": 0.2156008929014206,
+          "test_best_rerun_mse": null
+        },
+        "h_lo__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": 0.2164412885904312,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": 0.34223392605781555,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": 0.30678433179855347,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": 0.38425394892692566,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": 0.4129803776741028,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": 0.5420964360237122,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": 0.535067617893219,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": 0.6864542961120605,
+          "test_best_rerun_mse": null
+        }
+      },
+      "output_dir": "/raven/u/klettl/Projects/topobench/logs/train/runs/notebook_gu_grid_2026-07-29_18-09-01__community_detection__10__h_hi__d_hi__pl_lo__s43",
+      "wandb_config": {
+        "AvgTime/train_epoch_mean": 28.040712780422634,
+        "AvgTime/train_epoch_std": 1.5348362930251451,
+        "model/params/total": 682465,
+        "model/params/trainable": 682465,
+        "model/params/non_trainable": 0
+      }
+    },
+    {
+      "experiment": "community_detection",
+      "wandb_project": "challenge_community_detection",
+      "wandb_run_name": "gauge_hom_0.9-1__deg_4-5__gamma_1.5-2__s44",
+      "train_seed": 44,
+      "homophily": "h_hi",
+      "avg_degree": "d_hi",
+      "power_law": "pl_lo",
+      "run_slug": "h_hi__d_hi__pl_lo",
+      "test_loss": 1.4423375129699707,
+      "test_best_rerun_accuracy": 0.6401940584182739,
+      "test_best_rerun_mse": null,
+      "test_triangles_total_structural": null,
+      "test_mse_by_total_triangles": null,
+      "ood_test": {
+        "h_lo__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": 0.1776682734489441,
+          "test_best_rerun_mse": null
+        },
+        "h_lo__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": 0.16659027338027954,
+          "test_best_rerun_mse": null
+        },
+        "h_lo__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": 0.21460768580436707,
+          "test_best_rerun_mse": null
+        },
+        "h_lo__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": 0.22156009078025818,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": 0.3462449312210083,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": 0.3077393174171448,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": 0.3763083517551422,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": 0.41236916184425354,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": 0.5405302047729492,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": 0.5378562211990356,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": 0.681411862373352,
+          "test_best_rerun_mse": null
+        }
+      },
+      "output_dir": "/raven/u/klettl/Projects/topobench/logs/train/runs/notebook_gu_grid_2026-07-29_18-09-01__community_detection__10__h_hi__d_hi__pl_lo__s44",
+      "wandb_config": {
+        "AvgTime/train_epoch_mean": 27.487954223155974,
+        "AvgTime/train_epoch_std": 0.031968709742045305,
+        "model/params/total": 682465,
+        "model/params/trainable": 682465,
+        "model/params/non_trainable": 0
+      }
+    },
+    {
+      "experiment": "community_detection",
+      "wandb_project": "challenge_community_detection",
+      "wandb_run_name": "gauge_hom_0.9-1__deg_4-5__gamma_4-5__s42",
+      "train_seed": 42,
+      "homophily": "h_hi",
+      "avg_degree": "d_hi",
+      "power_law": "pl_hi",
+      "run_slug": "h_hi__d_hi__pl_hi",
+      "test_loss": 1.2890533208847046,
+      "test_best_rerun_accuracy": 0.6938268542289734,
+      "test_best_rerun_mse": null,
+      "test_triangles_total_structural": null,
+      "test_mse_by_total_triangles": null,
+      "ood_test": {
+        "h_lo__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": 0.174574077129364,
+          "test_best_rerun_mse": null
+        },
+        "h_lo__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": 0.16999006271362305,
+          "test_best_rerun_mse": null
+        },
+        "h_lo__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": 0.19088546931743622,
+          "test_best_rerun_mse": null
+        },
+        "h_lo__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": 0.19581328332424164,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": 0.3282909393310547,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": 0.30919092893600464,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": 0.3621743321418762,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": 0.4083581566810608,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": 0.5299488306045532,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": 0.5315532088279724,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": 0.6315990686416626,
+          "test_best_rerun_mse": null
+        }
+      },
+      "output_dir": "/raven/u/klettl/Projects/topobench/logs/train/runs/notebook_gu_grid_2026-07-29_18-09-01__community_detection__11__h_hi__d_hi__pl_hi__s42",
+      "wandb_config": {
+        "AvgTime/train_epoch_mean": 27.518502847353616,
+        "AvgTime/train_epoch_std": 0.04454948052478267,
+        "model/params/total": 682465,
+        "model/params/trainable": 682465,
+        "model/params/non_trainable": 0
+      }
+    },
+    {
+      "experiment": "community_detection",
+      "wandb_project": "challenge_community_detection",
+      "wandb_run_name": "gauge_hom_0.9-1__deg_4-5__gamma_4-5__s43",
+      "train_seed": 43,
+      "homophily": "h_hi",
+      "avg_degree": "d_hi",
+      "power_law": "pl_hi",
+      "run_slug": "h_hi__d_hi__pl_hi",
+      "test_loss": 1.3402941226959229,
+      "test_best_rerun_accuracy": 0.6873710751533508,
+      "test_best_rerun_mse": null,
+      "test_triangles_total_structural": null,
+      "test_mse_by_total_triangles": null,
+      "ood_test": {
+        "h_lo__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": 0.17625486850738525,
+          "test_best_rerun_mse": null
+        },
+        "h_lo__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": 0.16976086795330048,
+          "test_best_rerun_mse": null
+        },
+        "h_lo__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": 0.2047138810157776,
+          "test_best_rerun_mse": null
+        },
+        "h_lo__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": 0.20440828800201416,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": 0.3272213339805603,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": 0.30972573161125183,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": 0.382267564535141,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": 0.4147375524044037,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": 0.5244861841201782,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": 0.5268546342849731,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": 0.6270150542259216,
+          "test_best_rerun_mse": null
+        }
+      },
+      "output_dir": "/raven/u/klettl/Projects/topobench/logs/train/runs/notebook_gu_grid_2026-07-29_18-09-01__community_detection__11__h_hi__d_hi__pl_hi__s43",
+      "wandb_config": {
+        "AvgTime/train_epoch_mean": 27.37242741882801,
+        "AvgTime/train_epoch_std": 0.15180798189221864,
+        "model/params/total": 682465,
+        "model/params/trainable": 682465,
+        "model/params/non_trainable": 0
+      }
+    },
+    {
+      "experiment": "community_detection",
+      "wandb_project": "challenge_community_detection",
+      "wandb_run_name": "gauge_hom_0.9-1__deg_4-5__gamma_4-5__s44",
+      "train_seed": 44,
+      "homophily": "h_hi",
+      "avg_degree": "d_hi",
+      "power_law": "pl_hi",
+      "run_slug": "h_hi__d_hi__pl_hi",
+      "test_loss": 1.215087652206421,
+      "test_best_rerun_accuracy": 0.6951256990432739,
+      "test_best_rerun_mse": null,
+      "test_triangles_total_structural": null,
+      "test_mse_by_total_triangles": null,
+      "ood_test": {
+        "h_lo__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": 0.1829780787229538,
+          "test_best_rerun_mse": null
+        },
+        "h_lo__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": 0.1771334707736969,
+          "test_best_rerun_mse": null
+        },
+        "h_lo__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": 0.20158147811889648,
+          "test_best_rerun_mse": null
+        },
+        "h_lo__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": 0.2007792741060257,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": 0.3363129198551178,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": 0.31912294030189514,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": 0.36805716156959534,
+          "test_best_rerun_mse": null
+        },
+        "h_mid__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": 0.40831995010375977,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": 0.5302926301956177,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": 0.5385820269584656,
+          "test_best_rerun_mse": null
+        },
+        "h_hi__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": 0.6314462423324585,
+          "test_best_rerun_mse": null
+        }
+      },
+      "output_dir": "/raven/u/klettl/Projects/topobench/logs/train/runs/notebook_gu_grid_2026-07-29_18-09-01__community_detection__11__h_hi__d_hi__pl_hi__s44",
+      "wandb_config": {
+        "AvgTime/train_epoch_mean": 27.429591644377936,
+        "AvgTime/train_epoch_std": 0.2438520778978258,
+        "model/params/total": 682465,
+        "model/params/trainable": 682465,
+        "model/params/non_trainable": 0
+      }
+    },
+    {
+      "experiment": "triangle_counting",
+      "wandb_project": "challenge_triangle_counting",
+      "wandb_run_name": "gauge_hom_0-0.1__deg_1-2.5__gamma_1.5-2__s42",
+      "train_seed": 42,
+      "homophily": "h_lo",
+      "avg_degree": "d_lo",
+      "power_law": "pl_lo",
+      "run_slug": "h_lo__d_lo__pl_lo",
+      "test_loss": 102.22747039794922,
+      "test_best_rerun_accuracy": null,
+      "test_best_rerun_mse": 98.18943786621094,
+      "test_triangles_total_structural": 1388.0,
+      "test_mse_by_total_triangles": 0.07074166993242863,
+      "ood_test": {
+        "h_lo__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 32.82376480102539,
+          "test_triangles_total_structural": 190.0,
+          "test_mse_by_total_triangles": 0.17275665684750205
+        },
+        "h_lo__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 10553.2578125,
+          "test_triangles_total_structural": 13185.0,
+          "test_mse_by_total_triangles": 0.800398772279105
+        },
+        "h_lo__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 150.7338409423828,
+          "test_triangles_total_structural": 2967.0,
+          "test_mse_by_total_triangles": 0.05080345161522845
+        },
+        "h_mid__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 7375.33935546875,
+          "test_triangles_total_structural": 7485.0,
+          "test_mse_by_total_triangles": 0.9853492792877422
+        },
+        "h_mid__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 143.77902221679688,
+          "test_triangles_total_structural": 1158.0,
+          "test_mse_by_total_triangles": 0.12416150450500593
+        },
+        "h_mid__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 170180.953125,
+          "test_triangles_total_structural": 43064.0,
+          "test_mse_by_total_triangles": 3.9518148134172395
+        },
+        "h_mid__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 14311.4169921875,
+          "test_triangles_total_structural": 14262.0,
+          "test_mse_by_total_triangles": 1.0034649412556094
+        },
+        "h_hi__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 45934.15234375,
+          "test_triangles_total_structural": 19509.0,
+          "test_mse_by_total_triangles": 2.354510858770311
+        },
+        "h_hi__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 3593.04296875,
+          "test_triangles_total_structural": 5625.0,
+          "test_mse_by_total_triangles": 0.6387631944444444
+        },
+        "h_hi__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 751810.75,
+          "test_triangles_total_structural": 88403.0,
+          "test_mse_by_total_triangles": 8.504357883782225
+        },
+        "h_hi__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 258928.859375,
+          "test_triangles_total_structural": 60093.0,
+          "test_mse_by_total_triangles": 4.308802345947115
+        }
+      },
+      "output_dir": "/raven/u/klettl/Projects/topobench/logs/train/runs/notebook_gu_grid_2026-07-29_18-09-01__triangle_counting__00__h_lo__d_lo__pl_lo__s42",
+      "wandb_config": {
+        "AvgTime/train_epoch_mean": 26.222087343533833,
+        "AvgTime/train_epoch_std": 0.011195730950412638,
+        "model/params/total": 680014,
+        "model/params/trainable": 680014,
+        "model/params/non_trainable": 0
+      }
+    },
+    {
+      "experiment": "triangle_counting",
+      "wandb_project": "challenge_triangle_counting",
+      "wandb_run_name": "gauge_hom_0-0.1__deg_1-2.5__gamma_1.5-2__s43",
+      "train_seed": 43,
+      "homophily": "h_lo",
+      "avg_degree": "d_lo",
+      "power_law": "pl_lo",
+      "run_slug": "h_lo__d_lo__pl_lo",
+      "test_loss": 133.00172424316406,
+      "test_best_rerun_accuracy": null,
+      "test_best_rerun_mse": 128.4129180908203,
+      "test_triangles_total_structural": 1388.0,
+      "test_mse_by_total_triangles": 0.09251651159281002,
+      "ood_test": {
+        "h_lo__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 49.791866302490234,
+          "test_triangles_total_structural": 190.0,
+          "test_mse_by_total_triangles": 0.26206245422363283
+        },
+        "h_lo__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 11503.376953125,
+          "test_triangles_total_structural": 13185.0,
+          "test_mse_by_total_triangles": 0.8724593821103527
+        },
+        "h_lo__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 197.10946655273438,
+          "test_triangles_total_structural": 2967.0,
+          "test_mse_by_total_triangles": 0.06643392873364826
+        },
+        "h_mid__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 7426.18896484375,
+          "test_triangles_total_structural": 7485.0,
+          "test_mse_by_total_triangles": 0.9921428142743821
+        },
+        "h_mid__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 147.86944580078125,
+          "test_triangles_total_structural": 1158.0,
+          "test_mse_by_total_triangles": 0.12769382193504425
+        },
+        "h_mid__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 172973.046875,
+          "test_triangles_total_structural": 43064.0,
+          "test_mse_by_total_triangles": 4.016650726244659
+        },
+        "h_mid__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 14762.625,
+          "test_triangles_total_structural": 14262.0,
+          "test_mse_by_total_triangles": 1.0351020193521245
+        },
+        "h_hi__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 45726.91796875,
+          "test_triangles_total_structural": 19509.0,
+          "test_mse_by_total_triangles": 2.343888357616997
+        },
+        "h_hi__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 3527.21923828125,
+          "test_triangles_total_structural": 5625.0,
+          "test_mse_by_total_triangles": 0.6270611979166667
+        },
+        "h_hi__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 754070.8125,
+          "test_triangles_total_structural": 88403.0,
+          "test_mse_by_total_triangles": 8.52992333404975
+        },
+        "h_hi__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 260207.625,
+          "test_triangles_total_structural": 60093.0,
+          "test_mse_by_total_triangles": 4.3300821227098
+        }
+      },
+      "output_dir": "/raven/u/klettl/Projects/topobench/logs/train/runs/notebook_gu_grid_2026-07-29_18-09-01__triangle_counting__00__h_lo__d_lo__pl_lo__s43",
+      "wandb_config": {
+        "AvgTime/train_epoch_mean": 26.702368915081024,
+        "AvgTime/train_epoch_std": 0.007749227657102139,
+        "model/params/total": 680014,
+        "model/params/trainable": 680014,
+        "model/params/non_trainable": 0
+      }
+    },
+    {
+      "experiment": "triangle_counting",
+      "wandb_project": "challenge_triangle_counting",
+      "wandb_run_name": "gauge_hom_0-0.1__deg_1-2.5__gamma_1.5-2__s44",
+      "train_seed": 44,
+      "homophily": "h_lo",
+      "avg_degree": "d_lo",
+      "power_law": "pl_lo",
+      "run_slug": "h_lo__d_lo__pl_lo",
+      "test_loss": 108.23243713378906,
+      "test_best_rerun_accuracy": null,
+      "test_best_rerun_mse": 109.23623657226562,
+      "test_triangles_total_structural": 1388.0,
+      "test_mse_by_total_triangles": 0.0787004586255516,
+      "ood_test": {
+        "h_lo__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 26.19242286682129,
+          "test_triangles_total_structural": 190.0,
+          "test_mse_by_total_triangles": 0.13785485719379625
+        },
+        "h_lo__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 10168.236328125,
+          "test_triangles_total_structural": 13185.0,
+          "test_mse_by_total_triangles": 0.7711972945108078
+        },
+        "h_lo__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 134.8640594482422,
+          "test_triangles_total_structural": 2967.0,
+          "test_mse_by_total_triangles": 0.04545468805131182
+        },
+        "h_mid__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 7342.88330078125,
+          "test_triangles_total_structural": 7485.0,
+          "test_mse_by_total_triangles": 0.9810131330369072
+        },
+        "h_mid__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 146.48141479492188,
+          "test_triangles_total_structural": 1158.0,
+          "test_mse_by_total_triangles": 0.12649517685226414
+        },
+        "h_mid__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 169137.0625,
+          "test_triangles_total_structural": 43064.0,
+          "test_mse_by_total_triangles": 3.927574366059818
+        },
+        "h_mid__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 14136.0712890625,
+          "test_triangles_total_structural": 14262.0,
+          "test_mse_by_total_triangles": 0.9911703329871336
+        },
+        "h_hi__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 45871.09375,
+          "test_triangles_total_structural": 19509.0,
+          "test_mse_by_total_triangles": 2.3512785765544106
+        },
+        "h_hi__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 3602.771240234375,
+          "test_triangles_total_structural": 5625.0,
+          "test_mse_by_total_triangles": 0.6404926649305556
+        },
+        "h_hi__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 753496.1875,
+          "test_triangles_total_structural": 88403.0,
+          "test_mse_by_total_triangles": 8.523423271834666
+        },
+        "h_hi__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 258717.359375,
+          "test_triangles_total_structural": 60093.0,
+          "test_mse_by_total_triangles": 4.305282801241409
+        }
+      },
+      "output_dir": "/raven/u/klettl/Projects/topobench/logs/train/runs/notebook_gu_grid_2026-07-29_18-09-01__triangle_counting__00__h_lo__d_lo__pl_lo__s44",
+      "wandb_config": {
+        "AvgTime/train_epoch_mean": 26.64427587721083,
+        "AvgTime/train_epoch_std": 0.015150164592481927,
+        "model/params/total": 680014,
+        "model/params/trainable": 680014,
+        "model/params/non_trainable": 0
+      }
+    },
+    {
+      "experiment": "triangle_counting",
+      "wandb_project": "challenge_triangle_counting",
+      "wandb_run_name": "gauge_hom_0-0.1__deg_1-2.5__gamma_4-5__s42",
+      "train_seed": 42,
+      "homophily": "h_lo",
+      "avg_degree": "d_lo",
+      "power_law": "pl_hi",
+      "run_slug": "h_lo__d_lo__pl_hi",
+      "test_loss": 2.9345829486846924,
+      "test_best_rerun_accuracy": null,
+      "test_best_rerun_mse": 2.717726707458496,
+      "test_triangles_total_structural": 190.0,
+      "test_mse_by_total_triangles": 0.014303824776097348,
+      "ood_test": {
+        "h_lo__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 191.1142578125,
+          "test_triangles_total_structural": 1388.0,
+          "test_mse_by_total_triangles": 0.1376903874729827
+        },
+        "h_lo__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 12717.9287109375,
+          "test_triangles_total_structural": 13185.0,
+          "test_mse_by_total_triangles": 0.9645755563850967
+        },
+        "h_lo__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 358.98779296875,
+          "test_triangles_total_structural": 2967.0,
+          "test_mse_by_total_triangles": 0.12099352644716886
+        },
+        "h_mid__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 8079.7490234375,
+          "test_triangles_total_structural": 7485.0,
+          "test_mse_by_total_triangles": 1.0794587873663994
+        },
+        "h_mid__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 173.2749481201172,
+          "test_triangles_total_structural": 1158.0,
+          "test_mse_by_total_triangles": 0.14963294310890948
+        },
+        "h_mid__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 177214.828125,
+          "test_triangles_total_structural": 43064.0,
+          "test_mse_by_total_triangles": 4.115150197961174
+        },
+        "h_mid__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 15950.0498046875,
+          "test_triangles_total_structural": 14262.0,
+          "test_mse_by_total_triangles": 1.118359963868146
+        },
+        "h_hi__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 47582.83984375,
+          "test_triangles_total_structural": 19509.0,
+          "test_mse_by_total_triangles": 2.4390199315059715
+        },
+        "h_hi__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 3955.8837890625,
+          "test_triangles_total_structural": 5625.0,
+          "test_mse_by_total_triangles": 0.7032682291666666
+        },
+        "h_hi__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 763344.1875,
+          "test_triangles_total_structural": 88403.0,
+          "test_mse_by_total_triangles": 8.634822206259969
+        },
+        "h_hi__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 265947.96875,
+          "test_triangles_total_structural": 60093.0,
+          "test_mse_by_total_triangles": 4.425606455826802
+        }
+      },
+      "output_dir": "/raven/u/klettl/Projects/topobench/logs/train/runs/notebook_gu_grid_2026-07-29_18-09-01__triangle_counting__01__h_lo__d_lo__pl_hi__s42",
+      "wandb_config": {
+        "AvgTime/train_epoch_mean": 28.22666605710983,
+        "AvgTime/train_epoch_std": 1.6188551185071332,
+        "model/params/total": 680014,
+        "model/params/trainable": 680014,
+        "model/params/non_trainable": 0
+      }
+    },
+    {
+      "experiment": "triangle_counting",
+      "wandb_project": "challenge_triangle_counting",
+      "wandb_run_name": "gauge_hom_0-0.1__deg_1-2.5__gamma_4-5__s43",
+      "train_seed": 43,
+      "homophily": "h_lo",
+      "avg_degree": "d_lo",
+      "power_law": "pl_hi",
+      "run_slug": "h_lo__d_lo__pl_hi",
+      "test_loss": 2.5660552978515625,
+      "test_best_rerun_accuracy": null,
+      "test_best_rerun_mse": 2.371020793914795,
+      "test_triangles_total_structural": 190.0,
+      "test_mse_by_total_triangles": 0.012479056810077868,
+      "ood_test": {
+        "h_lo__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 199.7611541748047,
+          "test_triangles_total_structural": 1388.0,
+          "test_mse_by_total_triangles": 0.14392013989539243
+        },
+        "h_lo__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 12886.39453125,
+          "test_triangles_total_structural": 13185.0,
+          "test_mse_by_total_triangles": 0.9773526379408418
+        },
+        "h_lo__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 381.34375,
+          "test_triangles_total_structural": 2967.0,
+          "test_mse_by_total_triangles": 0.12852839568587798
+        },
+        "h_mid__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 8142.32080078125,
+          "test_triangles_total_structural": 7485.0,
+          "test_mse_by_total_triangles": 1.087818410258016
+        },
+        "h_mid__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 178.56500244140625,
+          "test_triangles_total_structural": 1158.0,
+          "test_mse_by_total_triangles": 0.15420121108929727
+        },
+        "h_mid__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 177841.1875,
+          "test_triangles_total_structural": 43064.0,
+          "test_mse_by_total_triangles": 4.129695046906929
+        },
+        "h_mid__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 16062.96484375,
+          "test_triangles_total_structural": 14262.0,
+          "test_mse_by_total_triangles": 1.1262771591466836
+        },
+        "h_hi__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 47769.62109375,
+          "test_triangles_total_structural": 19509.0,
+          "test_mse_by_total_triangles": 2.448594038328464
+        },
+        "h_hi__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 3980.1875,
+          "test_triangles_total_structural": 5625.0,
+          "test_mse_by_total_triangles": 0.7075888888888889
+        },
+        "h_hi__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 764871.5,
+          "test_triangles_total_structural": 88403.0,
+          "test_mse_by_total_triangles": 8.652098910670452
+        },
+        "h_hi__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 266582.8125,
+          "test_triangles_total_structural": 60093.0,
+          "test_mse_by_total_triangles": 4.436170810244121
+        }
+      },
+      "output_dir": "/raven/u/klettl/Projects/topobench/logs/train/runs/notebook_gu_grid_2026-07-29_18-09-01__triangle_counting__01__h_lo__d_lo__pl_hi__s43",
+      "wandb_config": {
+        "AvgTime/train_epoch_mean": 29.500447246763443,
+        "AvgTime/train_epoch_std": 1.1346180112479738,
+        "model/params/total": 680014,
+        "model/params/trainable": 680014,
+        "model/params/non_trainable": 0
+      }
+    },
+    {
+      "experiment": "triangle_counting",
+      "wandb_project": "challenge_triangle_counting",
+      "wandb_run_name": "gauge_hom_0-0.1__deg_1-2.5__gamma_4-5__s44",
+      "train_seed": 44,
+      "homophily": "h_lo",
+      "avg_degree": "d_lo",
+      "power_law": "pl_hi",
+      "run_slug": "h_lo__d_lo__pl_hi",
+      "test_loss": 2.9012434482574463,
+      "test_best_rerun_accuracy": null,
+      "test_best_rerun_mse": 2.6825222969055176,
+      "test_triangles_total_structural": 190.0,
+      "test_mse_by_total_triangles": 0.014118538404765882,
+      "ood_test": {
+        "h_lo__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 207.81871032714844,
+          "test_triangles_total_structural": 1388.0,
+          "test_mse_by_total_triangles": 0.1497252956247467
+        },
+        "h_lo__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 13124.4228515625,
+          "test_triangles_total_structural": 13185.0,
+          "test_mse_by_total_triangles": 0.9954056011803185
+        },
+        "h_lo__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 416.5335693359375,
+          "test_triangles_total_structural": 2967.0,
+          "test_mse_by_total_triangles": 0.14038879991100017
+        },
+        "h_mid__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 8170.22412109375,
+          "test_triangles_total_structural": 7485.0,
+          "test_mse_by_total_triangles": 1.0915463087633601
+        },
+        "h_mid__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 175.6750946044922,
+          "test_triangles_total_structural": 1158.0,
+          "test_mse_by_total_triangles": 0.15170560846674627
+        },
+        "h_mid__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 178562.65625,
+          "test_triangles_total_structural": 43064.0,
+          "test_mse_by_total_triangles": 4.1464484546256735
+        },
+        "h_mid__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 16227.5263671875,
+          "test_triangles_total_structural": 14262.0,
+          "test_mse_by_total_triangles": 1.1378156196317137
+        },
+        "h_hi__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 47822.09765625,
+          "test_triangles_total_structural": 19509.0,
+          "test_mse_by_total_triangles": 2.4512839026218667
+        },
+        "h_hi__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 3984.58203125,
+          "test_triangles_total_structural": 5625.0,
+          "test_mse_by_total_triangles": 0.7083701388888889
+        },
+        "h_hi__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 765973.8125,
+          "test_triangles_total_structural": 88403.0,
+          "test_mse_by_total_triangles": 8.664568085924685
+        },
+        "h_hi__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 267091.78125,
+          "test_triangles_total_structural": 60093.0,
+          "test_mse_by_total_triangles": 4.4446404947331635
+        }
+      },
+      "output_dir": "/raven/u/klettl/Projects/topobench/logs/train/runs/notebook_gu_grid_2026-07-29_18-09-01__triangle_counting__01__h_lo__d_lo__pl_hi__s44",
+      "wandb_config": {
+        "AvgTime/train_epoch_mean": 26.313814401626587,
+        "AvgTime/train_epoch_std": 0.016592741012573242,
+        "model/params/total": 680014,
+        "model/params/trainable": 680014,
+        "model/params/non_trainable": 0
+      }
+    },
+    {
+      "experiment": "triangle_counting",
+      "wandb_project": "challenge_triangle_counting",
+      "wandb_run_name": "gauge_hom_0-0.1__deg_4-5__gamma_1.5-2__s42",
+      "train_seed": 42,
+      "homophily": "h_lo",
+      "avg_degree": "d_hi",
+      "power_law": "pl_lo",
+      "run_slug": "h_lo__d_hi__pl_lo",
+      "test_loss": 4824.7412109375,
+      "test_best_rerun_accuracy": null,
+      "test_best_rerun_mse": 4562.404296875,
+      "test_triangles_total_structural": 13185.0,
+      "test_mse_by_total_triangles": 0.34602990495828595,
+      "ood_test": {
+        "h_lo__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 5783.71240234375,
+          "test_triangles_total_structural": 1388.0,
+          "test_mse_by_total_triangles": 4.16693977114103
+        },
+        "h_lo__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 7076.76318359375,
+          "test_triangles_total_structural": 190.0,
+          "test_mse_by_total_triangles": 37.24612201891448
+        },
+        "h_lo__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 4950.18212890625,
+          "test_triangles_total_structural": 2967.0,
+          "test_mse_by_total_triangles": 1.668413255445315
+        },
+        "h_mid__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 6258.08935546875,
+          "test_triangles_total_structural": 7485.0,
+          "test_mse_by_total_triangles": 0.836084082226954
+        },
+        "h_mid__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 6244.35546875,
+          "test_triangles_total_structural": 1158.0,
+          "test_mse_by_total_triangles": 5.392362235535406
+        },
+        "h_mid__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 130713.8671875,
+          "test_triangles_total_structural": 43064.0,
+          "test_mse_by_total_triangles": 3.0353396616083037
+        },
+        "h_mid__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 8247.0859375,
+          "test_triangles_total_structural": 14262.0,
+          "test_mse_by_total_triangles": 0.5782559204529519
+        },
+        "h_hi__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 30823.32421875,
+          "test_triangles_total_structural": 19509.0,
+          "test_mse_by_total_triangles": 1.5799540836921422
+        },
+        "h_hi__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 5235.18994140625,
+          "test_triangles_total_structural": 5625.0,
+          "test_mse_by_total_triangles": 0.9307004340277778
+        },
+        "h_hi__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 657184.6875,
+          "test_triangles_total_structural": 88403.0,
+          "test_mse_by_total_triangles": 7.433963638111829
+        },
+        "h_hi__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 205770.21875,
+          "test_triangles_total_structural": 60093.0,
+          "test_mse_by_total_triangles": 3.4241961418135225
+        }
+      },
+      "output_dir": "/raven/u/klettl/Projects/topobench/logs/train/runs/notebook_gu_grid_2026-07-29_18-09-01__triangle_counting__02__h_lo__d_hi__pl_lo__s42",
+      "wandb_config": {
+        "AvgTime/train_epoch_mean": 26.873037338256836,
+        "AvgTime/train_epoch_std": 0.09505019037719786,
+        "model/params/total": 680014,
+        "model/params/trainable": 680014,
+        "model/params/non_trainable": 0
+      }
+    },
+    {
+      "experiment": "triangle_counting",
+      "wandb_project": "challenge_triangle_counting",
+      "wandb_run_name": "gauge_hom_0-0.1__deg_4-5__gamma_1.5-2__s43",
+      "train_seed": 43,
+      "homophily": "h_lo",
+      "avg_degree": "d_hi",
+      "power_law": "pl_lo",
+      "run_slug": "h_lo__d_hi__pl_lo",
+      "test_loss": 4903.6416015625,
+      "test_best_rerun_accuracy": null,
+      "test_best_rerun_mse": 4619.92578125,
+      "test_triangles_total_structural": 13185.0,
+      "test_mse_by_total_triangles": 0.3503925507205157,
+      "ood_test": {
+        "h_lo__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 6648.6728515625,
+          "test_triangles_total_structural": 1388.0,
+          "test_mse_by_total_triangles": 4.790110123604107
+        },
+        "h_lo__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 8010.90478515625,
+          "test_triangles_total_structural": 190.0,
+          "test_mse_by_total_triangles": 42.162656763980266
+        },
+        "h_lo__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 5760.06689453125,
+          "test_triangles_total_structural": 2967.0,
+          "test_mse_by_total_triangles": 1.9413774501284968
+        },
+        "h_mid__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 6645.85400390625,
+          "test_triangles_total_structural": 7485.0,
+          "test_mse_by_total_triangles": 0.8878896464804609
+        },
+        "h_mid__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 7074.61083984375,
+          "test_triangles_total_structural": 1158.0,
+          "test_mse_by_total_triangles": 6.10933578570272
+        },
+        "h_mid__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 128502.53125,
+          "test_triangles_total_structural": 43064.0,
+          "test_mse_by_total_triangles": 2.983989672348133
+        },
+        "h_mid__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 8231.263671875,
+          "test_triangles_total_structural": 14262.0,
+          "test_mse_by_total_triangles": 0.5771465202548731
+        },
+        "h_hi__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 30276.3203125,
+          "test_triangles_total_structural": 19509.0,
+          "test_mse_by_total_triangles": 1.5519155421856579
+        },
+        "h_hi__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 5665.05810546875,
+          "test_triangles_total_structural": 5625.0,
+          "test_mse_by_total_triangles": 1.0071214409722222
+        },
+        "h_hi__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 653691.6875,
+          "test_triangles_total_structural": 88403.0,
+          "test_mse_by_total_triangles": 7.394451404364105
+        },
+        "h_hi__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 203101.671875,
+          "test_triangles_total_structural": 60093.0,
+          "test_mse_by_total_triangles": 3.3797891913367613
+        }
+      },
+      "output_dir": "/raven/u/klettl/Projects/topobench/logs/train/runs/notebook_gu_grid_2026-07-29_18-09-01__triangle_counting__02__h_lo__d_hi__pl_lo__s43",
+      "wandb_config": {
+        "AvgTime/train_epoch_mean": 30.212888717651367,
+        "AvgTime/train_epoch_std": 0.015450738853387387,
+        "model/params/total": 680014,
+        "model/params/trainable": 680014,
+        "model/params/non_trainable": 0
+      }
+    },
+    {
+      "experiment": "triangle_counting",
+      "wandb_project": "challenge_triangle_counting",
+      "wandb_run_name": "gauge_hom_0-0.1__deg_4-5__gamma_1.5-2__s44",
+      "train_seed": 44,
+      "homophily": "h_lo",
+      "avg_degree": "d_hi",
+      "power_law": "pl_lo",
+      "run_slug": "h_lo__d_hi__pl_lo",
+      "test_loss": 5000.74072265625,
+      "test_best_rerun_accuracy": null,
+      "test_best_rerun_mse": 4682.89990234375,
+      "test_triangles_total_structural": 13185.0,
+      "test_mse_by_total_triangles": 0.3551687449635002,
+      "ood_test": {
+        "h_lo__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 6078.462890625,
+          "test_triangles_total_structural": 1388.0,
+          "test_mse_by_total_triangles": 4.3792960307096545
+        },
+        "h_lo__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 7414.9814453125,
+          "test_triangles_total_structural": 190.0,
+          "test_mse_by_total_triangles": 39.026218133223686
+        },
+        "h_lo__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 5140.5439453125,
+          "test_triangles_total_structural": 2967.0,
+          "test_mse_by_total_triangles": 1.732572950897371
+        },
+        "h_mid__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 6437.06494140625,
+          "test_triangles_total_structural": 7485.0,
+          "test_mse_by_total_triangles": 0.8599953161531396
+        },
+        "h_mid__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 6599.9482421875,
+          "test_triangles_total_structural": 1158.0,
+          "test_mse_by_total_triangles": 5.699437169419257
+        },
+        "h_mid__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 130551.8671875,
+          "test_triangles_total_structural": 43064.0,
+          "test_mse_by_total_triangles": 3.031577818769738
+        },
+        "h_mid__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 8187.146484375,
+          "test_triangles_total_structural": 14262.0,
+          "test_mse_by_total_triangles": 0.5740531821886832
+        },
+        "h_hi__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 30665.953125,
+          "test_triangles_total_structural": 19509.0,
+          "test_mse_by_total_triangles": 1.5718874942334307
+        },
+        "h_hi__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 5355.583984375,
+          "test_triangles_total_structural": 5625.0,
+          "test_mse_by_total_triangles": 0.9521038194444444
+        },
+        "h_hi__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 657362.9375,
+          "test_triangles_total_structural": 88403.0,
+          "test_mse_by_total_triangles": 7.435979972399127
+        },
+        "h_hi__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 205180.5625,
+          "test_triangles_total_structural": 60093.0,
+          "test_mse_by_total_triangles": 3.414383746859035
+        }
+      },
+      "output_dir": "/raven/u/klettl/Projects/topobench/logs/train/runs/notebook_gu_grid_2026-07-29_18-09-01__triangle_counting__02__h_lo__d_hi__pl_lo__s44",
+      "wandb_config": {
+        "AvgTime/train_epoch_mean": 26.81612801551819,
+        "AvgTime/train_epoch_std": 0.013562202453613281,
+        "model/params/total": 680014,
+        "model/params/trainable": 680014,
+        "model/params/non_trainable": 0
+      }
+    },
+    {
+      "experiment": "triangle_counting",
+      "wandb_project": "challenge_triangle_counting",
+      "wandb_run_name": "gauge_hom_0-0.1__deg_4-5__gamma_4-5__s42",
+      "train_seed": 42,
+      "homophily": "h_lo",
+      "avg_degree": "d_hi",
+      "power_law": "pl_hi",
+      "run_slug": "h_lo__d_hi__pl_hi",
+      "test_loss": 138.8093719482422,
+      "test_best_rerun_accuracy": null,
+      "test_best_rerun_mse": 134.63877868652344,
+      "test_triangles_total_structural": 2967.0,
+      "test_mse_by_total_triangles": 0.04537875924722731,
+      "ood_test": {
+        "h_lo__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 170.32859802246094,
+          "test_triangles_total_structural": 1388.0,
+          "test_mse_by_total_triangles": 0.12271512825825716
+        },
+        "h_lo__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 224.63551330566406,
+          "test_triangles_total_structural": 190.0,
+          "test_mse_by_total_triangles": 1.1822921752929687
+        },
+        "h_lo__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 10132.875,
+          "test_triangles_total_structural": 13185.0,
+          "test_mse_by_total_triangles": 0.7685153583617748
+        },
+        "h_mid__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 6757.9794921875,
+          "test_triangles_total_structural": 7485.0,
+          "test_mse_by_total_triangles": 0.9028696716349366
+        },
+        "h_mid__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 224.60498046875,
+          "test_triangles_total_structural": 1158.0,
+          "test_mse_by_total_triangles": 0.19395939591429187
+        },
+        "h_mid__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 167627.03125,
+          "test_triangles_total_structural": 43064.0,
+          "test_mse_by_total_triangles": 3.892509549739922
+        },
+        "h_mid__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 13474.4423828125,
+          "test_triangles_total_structural": 14262.0,
+          "test_mse_by_total_triangles": 0.9447793004355981
+        },
+        "h_hi__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 43835.83203125,
+          "test_triangles_total_structural": 19509.0,
+          "test_mse_by_total_triangles": 2.246954330373161
+        },
+        "h_hi__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 3168.111572265625,
+          "test_triangles_total_structural": 5625.0,
+          "test_mse_by_total_triangles": 0.5632198350694444
+        },
+        "h_hi__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 744902.1875,
+          "test_triangles_total_structural": 88403.0,
+          "test_mse_by_total_triangles": 8.42620937637863
+        },
+        "h_hi__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 254407.859375,
+          "test_triangles_total_structural": 60093.0,
+          "test_mse_by_total_triangles": 4.2335689576989
+        }
+      },
+      "output_dir": "/raven/u/klettl/Projects/topobench/logs/train/runs/notebook_gu_grid_2026-07-29_18-09-01__triangle_counting__03__h_lo__d_hi__pl_hi__s42",
+      "wandb_config": {
+        "AvgTime/train_epoch_mean": 26.902429035731725,
+        "AvgTime/train_epoch_std": 0.013756220402894209,
+        "model/params/total": 680014,
+        "model/params/trainable": 680014,
+        "model/params/non_trainable": 0
+      }
+    },
+    {
+      "experiment": "triangle_counting",
+      "wandb_project": "challenge_triangle_counting",
+      "wandb_run_name": "gauge_hom_0-0.1__deg_4-5__gamma_4-5__s43",
+      "train_seed": 43,
+      "homophily": "h_lo",
+      "avg_degree": "d_hi",
+      "power_law": "pl_hi",
+      "run_slug": "h_lo__d_hi__pl_hi",
+      "test_loss": 133.72964477539062,
+      "test_best_rerun_accuracy": null,
+      "test_best_rerun_mse": 130.73651123046875,
+      "test_triangles_total_structural": 2967.0,
+      "test_mse_by_total_triangles": 0.04406353597252064,
+      "ood_test": {
+        "h_lo__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 175.15072631835938,
+          "test_triangles_total_structural": 1388.0,
+          "test_mse_by_total_triangles": 0.12618928409103702
+        },
+        "h_lo__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 247.11376953125,
+          "test_triangles_total_structural": 190.0,
+          "test_mse_by_total_triangles": 1.300598787006579
+        },
+        "h_lo__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 10034.2919921875,
+          "test_triangles_total_structural": 13185.0,
+          "test_mse_by_total_triangles": 0.7610384521947289
+        },
+        "h_mid__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 6731.8408203125,
+          "test_triangles_total_structural": 7485.0,
+          "test_mse_by_total_triangles": 0.8993775311038744
+        },
+        "h_mid__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 236.10726928710938,
+          "test_triangles_total_structural": 1158.0,
+          "test_mse_by_total_triangles": 0.20389228781270238
+        },
+        "h_mid__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 167620.453125,
+          "test_triangles_total_structural": 43064.0,
+          "test_mse_by_total_triangles": 3.892356797441018
+        },
+        "h_mid__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 13428.392578125,
+          "test_triangles_total_structural": 14262.0,
+          "test_mse_by_total_triangles": 0.9415504542227597
+        },
+        "h_hi__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 43603.96484375,
+          "test_triangles_total_structural": 19509.0,
+          "test_mse_by_total_triangles": 2.2350691908221845
+        },
+        "h_hi__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 3151.350341796875,
+          "test_triangles_total_structural": 5625.0,
+          "test_mse_by_total_triangles": 0.5602400607638889
+        },
+        "h_hi__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 744044.0,
+          "test_triangles_total_structural": 88403.0,
+          "test_mse_by_total_triangles": 8.416501702430914
+        },
+        "h_hi__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 253894.03125,
+          "test_triangles_total_structural": 60093.0,
+          "test_mse_by_total_triangles": 4.225018408966102
+        }
+      },
+      "output_dir": "/raven/u/klettl/Projects/topobench/logs/train/runs/notebook_gu_grid_2026-07-29_18-09-01__triangle_counting__03__h_lo__d_hi__pl_hi__s43",
+      "wandb_config": {
+        "AvgTime/train_epoch_mean": 26.99335026741028,
+        "AvgTime/train_epoch_std": 0.03801027699307579,
+        "model/params/total": 680014,
+        "model/params/trainable": 680014,
+        "model/params/non_trainable": 0
+      }
+    },
+    {
+      "experiment": "triangle_counting",
+      "wandb_project": "challenge_triangle_counting",
+      "wandb_run_name": "gauge_hom_0-0.1__deg_4-5__gamma_4-5__s44",
+      "train_seed": 44,
+      "homophily": "h_lo",
+      "avg_degree": "d_hi",
+      "power_law": "pl_hi",
+      "run_slug": "h_lo__d_hi__pl_hi",
+      "test_loss": 140.10427856445312,
+      "test_best_rerun_accuracy": null,
+      "test_best_rerun_mse": 136.88597106933594,
+      "test_triangles_total_structural": 2967.0,
+      "test_mse_by_total_triangles": 0.04613615472508795,
+      "ood_test": {
+        "h_lo__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 214.4286346435547,
+          "test_triangles_total_structural": 1388.0,
+          "test_mse_by_total_triangles": 0.1544874889362786
+        },
+        "h_lo__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 321.9037780761719,
+          "test_triangles_total_structural": 190.0,
+          "test_mse_by_total_triangles": 1.6942304109272204
+        },
+        "h_lo__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 10164.412109375,
+          "test_triangles_total_structural": 13185.0,
+          "test_mse_by_total_triangles": 0.7709072513746682
+        },
+        "h_mid__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 6612.8388671875,
+          "test_triangles_total_structural": 7485.0,
+          "test_mse_by_total_triangles": 0.8834788065714763
+        },
+        "h_mid__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 292.9797058105469,
+          "test_triangles_total_structural": 1158.0,
+          "test_mse_by_total_triangles": 0.25300492729753615
+        },
+        "h_mid__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 167357.140625,
+          "test_triangles_total_structural": 43064.0,
+          "test_mse_by_total_triangles": 3.886242351500093
+        },
+        "h_mid__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 13309.4794921875,
+          "test_triangles_total_structural": 14262.0,
+          "test_mse_by_total_triangles": 0.9332126975310265
+        },
+        "h_hi__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 42999.44921875,
+          "test_triangles_total_structural": 19509.0,
+          "test_mse_by_total_triangles": 2.204082691001589
+        },
+        "h_hi__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 3047.14892578125,
+          "test_triangles_total_structural": 5625.0,
+          "test_mse_by_total_triangles": 0.5417153645833334
+        },
+        "h_hi__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 742018.8125,
+          "test_triangles_total_structural": 88403.0,
+          "test_mse_by_total_triangles": 8.393593119011799
+        },
+        "h_hi__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 252939.734375,
+          "test_triangles_total_structural": 60093.0,
+          "test_mse_by_total_triangles": 4.209138075566206
+        }
+      },
+      "output_dir": "/raven/u/klettl/Projects/topobench/logs/train/runs/notebook_gu_grid_2026-07-29_18-09-01__triangle_counting__03__h_lo__d_hi__pl_hi__s44",
+      "wandb_config": {
+        "AvgTime/train_epoch_mean": 26.89208745956421,
+        "AvgTime/train_epoch_std": 0.01713395118713379,
+        "model/params/total": 680014,
+        "model/params/trainable": 680014,
+        "model/params/non_trainable": 0
+      }
+    },
+    {
+      "experiment": "triangle_counting",
+      "wandb_project": "challenge_triangle_counting",
+      "wandb_run_name": "gauge_hom_0.4-0.6__deg_1-2.5__gamma_1.5-2__s42",
+      "train_seed": 42,
+      "homophily": "h_mid",
+      "avg_degree": "d_lo",
+      "power_law": "pl_lo",
+      "run_slug": "h_mid__d_lo__pl_lo",
+      "test_loss": 3487.518798828125,
+      "test_best_rerun_accuracy": null,
+      "test_best_rerun_mse": 3534.981201171875,
+      "test_triangles_total_structural": 7485.0,
+      "test_mse_by_total_triangles": 0.4722753775780728,
+      "ood_test": {
+        "h_lo__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 879.0169677734375,
+          "test_triangles_total_structural": 1388.0,
+          "test_mse_by_total_triangles": 0.6332975272142921
+        },
+        "h_lo__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 439.8455810546875,
+          "test_triangles_total_structural": 190.0,
+          "test_mse_by_total_triangles": 2.314976742393092
+        },
+        "h_lo__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 4756.21435546875,
+          "test_triangles_total_structural": 13185.0,
+          "test_mse_by_total_triangles": 0.3607291888865188
+        },
+        "h_lo__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 7325.38427734375,
+          "test_triangles_total_structural": 2967.0,
+          "test_mse_by_total_triangles": 2.4689532448074654
+        },
+        "h_mid__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 511.31427001953125,
+          "test_triangles_total_structural": 1158.0,
+          "test_mse_by_total_triangles": 0.44154945597541556
+        },
+        "h_mid__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 96317.453125,
+          "test_triangles_total_structural": 43064.0,
+          "test_mse_by_total_triangles": 2.2366118596739737
+        },
+        "h_mid__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 7403.35400390625,
+          "test_triangles_total_structural": 14262.0,
+          "test_mse_by_total_triangles": 0.5190964804309529
+        },
+        "h_hi__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 33821.390625,
+          "test_triangles_total_structural": 19509.0,
+          "test_mse_by_total_triangles": 1.733630151468553
+        },
+        "h_hi__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 2750.083984375,
+          "test_triangles_total_structural": 5625.0,
+          "test_mse_by_total_triangles": 0.48890381944444444
+        },
+        "h_hi__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 631555.4375,
+          "test_triangles_total_structural": 88403.0,
+          "test_mse_by_total_triangles": 7.144049834281642
+        },
+        "h_hi__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 187862.890625,
+          "test_triangles_total_structural": 60093.0,
+          "test_mse_by_total_triangles": 3.126202563110512
+        }
+      },
+      "output_dir": "/raven/u/klettl/Projects/topobench/logs/train/runs/notebook_gu_grid_2026-07-29_18-09-01__triangle_counting__04__h_mid__d_lo__pl_lo__s42",
+      "wandb_config": {
+        "AvgTime/train_epoch_mean": 26.566348135471344,
+        "AvgTime/train_epoch_std": 0.1508598448936571,
+        "model/params/total": 680014,
+        "model/params/trainable": 680014,
+        "model/params/non_trainable": 0
+      }
+    },
+    {
+      "experiment": "triangle_counting",
+      "wandb_project": "challenge_triangle_counting",
+      "wandb_run_name": "gauge_hom_0.4-0.6__deg_1-2.5__gamma_1.5-2__s43",
+      "train_seed": 43,
+      "homophily": "h_mid",
+      "avg_degree": "d_lo",
+      "power_law": "pl_lo",
+      "run_slug": "h_mid__d_lo__pl_lo",
+      "test_loss": 4723.4677734375,
+      "test_best_rerun_accuracy": null,
+      "test_best_rerun_mse": 4753.767578125,
+      "test_triangles_total_structural": 7485.0,
+      "test_mse_by_total_triangles": 0.6351058888610555,
+      "ood_test": {
+        "h_lo__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 1049.284912109375,
+          "test_triangles_total_structural": 1388.0,
+          "test_mse_by_total_triangles": 0.7559689568511347
+        },
+        "h_lo__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 1029.009033203125,
+          "test_triangles_total_structural": 190.0,
+          "test_mse_by_total_triangles": 5.415837016858553
+        },
+        "h_lo__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 4612.619140625,
+          "test_triangles_total_structural": 13185.0,
+          "test_mse_by_total_triangles": 0.3498383876090254
+        },
+        "h_lo__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 2640.49658203125,
+          "test_triangles_total_structural": 2967.0,
+          "test_mse_by_total_triangles": 0.8899550327034884
+        },
+        "h_mid__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 904.728515625,
+          "test_triangles_total_structural": 1158.0,
+          "test_mse_by_total_triangles": 0.781285419365285
+        },
+        "h_mid__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 135458.375,
+          "test_triangles_total_structural": 43064.0,
+          "test_mse_by_total_triangles": 3.145513073564927
+        },
+        "h_mid__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 8513.7216796875,
+          "test_triangles_total_structural": 14262.0,
+          "test_mse_by_total_triangles": 0.5969514569967396
+        },
+        "h_hi__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 37149.296875,
+          "test_triangles_total_structural": 19509.0,
+          "test_mse_by_total_triangles": 1.904213279768312
+        },
+        "h_hi__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 2489.83203125,
+          "test_triangles_total_structural": 5625.0,
+          "test_mse_by_total_triangles": 0.44263680555555557
+        },
+        "h_hi__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 694613.4375,
+          "test_triangles_total_structural": 88403.0,
+          "test_mse_by_total_triangles": 7.857351419069489
+        },
+        "h_hi__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 221754.84375,
+          "test_triangles_total_structural": 60093.0,
+          "test_mse_by_total_triangles": 3.690194261394838
+        }
+      },
+      "output_dir": "/raven/u/klettl/Projects/topobench/logs/train/runs/notebook_gu_grid_2026-07-29_18-09-01__triangle_counting__04__h_mid__d_lo__pl_lo__s43",
+      "wandb_config": {
+        "AvgTime/train_epoch_mean": 26.343173674174718,
+        "AvgTime/train_epoch_std": 0.046057218900491595,
+        "model/params/total": 680014,
+        "model/params/trainable": 680014,
+        "model/params/non_trainable": 0
+      }
+    },
+    {
+      "experiment": "triangle_counting",
+      "wandb_project": "challenge_triangle_counting",
+      "wandb_run_name": "gauge_hom_0.4-0.6__deg_1-2.5__gamma_1.5-2__s44",
+      "train_seed": 44,
+      "homophily": "h_mid",
+      "avg_degree": "d_lo",
+      "power_law": "pl_lo",
+      "run_slug": "h_mid__d_lo__pl_lo",
+      "test_loss": 5225.74853515625,
+      "test_best_rerun_accuracy": null,
+      "test_best_rerun_mse": 5399.06689453125,
+      "test_triangles_total_structural": 7485.0,
+      "test_mse_by_total_triangles": 0.7213182223822645,
+      "ood_test": {
+        "h_lo__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 1403.268310546875,
+          "test_triangles_total_structural": 1388.0,
+          "test_mse_by_total_triangles": 1.0110002237369415
+        },
+        "h_lo__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 1973.6402587890625,
+          "test_triangles_total_structural": 190.0,
+          "test_mse_by_total_triangles": 10.38758030941612
+        },
+        "h_lo__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 6720.40087890625,
+          "test_triangles_total_structural": 13185.0,
+          "test_mse_by_total_triangles": 0.5097004838002465
+        },
+        "h_lo__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 970.983642578125,
+          "test_triangles_total_structural": 2967.0,
+          "test_mse_by_total_triangles": 0.3272610861402511
+        },
+        "h_mid__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 1636.0723876953125,
+          "test_triangles_total_structural": 1158.0,
+          "test_mse_by_total_triangles": 1.4128431672671093
+        },
+        "h_mid__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 151060.53125,
+          "test_triangles_total_structural": 43064.0,
+          "test_mse_by_total_triangles": 3.5078146769923833
+        },
+        "h_mid__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 10009.076171875,
+          "test_triangles_total_structural": 14262.0,
+          "test_mse_by_total_triangles": 0.7018003205633853
+        },
+        "h_hi__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 37087.30859375,
+          "test_triangles_total_structural": 19509.0,
+          "test_mse_by_total_triangles": 1.901035860051771
+        },
+        "h_hi__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 2837.887451171875,
+          "test_triangles_total_structural": 5625.0,
+          "test_mse_by_total_triangles": 0.5045133246527778
+        },
+        "h_hi__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 707834.6875,
+          "test_triangles_total_structural": 88403.0,
+          "test_mse_by_total_triangles": 8.006907995203782
+        },
+        "h_hi__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 232866.609375,
+          "test_triangles_total_structural": 60093.0,
+          "test_mse_by_total_triangles": 3.875103745444561
+        }
+      },
+      "output_dir": "/raven/u/klettl/Projects/topobench/logs/train/runs/notebook_gu_grid_2026-07-29_18-09-01__triangle_counting__04__h_mid__d_lo__pl_lo__s44",
+      "wandb_config": {
+        "AvgTime/train_epoch_mean": 26.320187211036682,
+        "AvgTime/train_epoch_std": 0.01039135456085205,
+        "model/params/total": 680014,
+        "model/params/trainable": 680014,
+        "model/params/non_trainable": 0
+      }
+    },
+    {
+      "experiment": "triangle_counting",
+      "wandb_project": "challenge_triangle_counting",
+      "wandb_run_name": "gauge_hom_0.4-0.6__deg_1-2.5__gamma_4-5__s42",
+      "train_seed": 42,
+      "homophily": "h_mid",
+      "avg_degree": "d_lo",
+      "power_law": "pl_hi",
+      "run_slug": "h_mid__d_lo__pl_hi",
+      "test_loss": 134.52127075195312,
+      "test_best_rerun_accuracy": null,
+      "test_best_rerun_mse": 141.66932678222656,
+      "test_triangles_total_structural": 1158.0,
+      "test_mse_by_total_triangles": 0.12233966043370169,
+      "ood_test": {
+        "h_lo__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 102.86193084716797,
+          "test_triangles_total_structural": 1388.0,
+          "test_mse_by_total_triangles": 0.07410801934234003
+        },
+        "h_lo__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 46.3376579284668,
+          "test_triangles_total_structural": 190.0,
+          "test_mse_by_total_triangles": 0.24388241014982526
+        },
+        "h_lo__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 10231.548828125,
+          "test_triangles_total_structural": 13185.0,
+          "test_mse_by_total_triangles": 0.775999152682973
+        },
+        "h_lo__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 136.13406372070312,
+          "test_triangles_total_structural": 2967.0,
+          "test_mse_by_total_triangles": 0.04588273128436236
+        },
+        "h_mid__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 7269.0224609375,
+          "test_triangles_total_structural": 7485.0,
+          "test_mse_by_total_triangles": 0.9711452853623914
+        },
+        "h_mid__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 169029.71875,
+          "test_triangles_total_structural": 43064.0,
+          "test_mse_by_total_triangles": 3.9250817097807915
+        },
+        "h_mid__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 13984.7529296875,
+          "test_triangles_total_structural": 14262.0,
+          "test_mse_by_total_triangles": 0.9805604354008904
+        },
+        "h_hi__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 45834.72265625,
+          "test_triangles_total_structural": 19509.0,
+          "test_mse_by_total_triangles": 2.349414252716695
+        },
+        "h_hi__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 3555.60791015625,
+          "test_triangles_total_structural": 5625.0,
+          "test_mse_by_total_triangles": 0.6321080729166667
+        },
+        "h_hi__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 751324.9375,
+          "test_triangles_total_structural": 88403.0,
+          "test_mse_by_total_triangles": 8.49886245376288
+        },
+        "h_hi__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 258070.109375,
+          "test_triangles_total_structural": 60093.0,
+          "test_mse_by_total_triangles": 4.2945119959895495
+        }
+      },
+      "output_dir": "/raven/u/klettl/Projects/topobench/logs/train/runs/notebook_gu_grid_2026-07-29_18-09-01__triangle_counting__05__h_mid__d_lo__pl_hi__s42",
+      "wandb_config": {
+        "AvgTime/train_epoch_mean": 26.832415014505386,
+        "AvgTime/train_epoch_std": 0.017686286370130776,
+        "model/params/total": 680014,
+        "model/params/trainable": 680014,
+        "model/params/non_trainable": 0
+      }
+    },
+    {
+      "experiment": "triangle_counting",
+      "wandb_project": "challenge_triangle_counting",
+      "wandb_run_name": "gauge_hom_0.4-0.6__deg_1-2.5__gamma_4-5__s43",
+      "train_seed": 43,
+      "homophily": "h_mid",
+      "avg_degree": "d_lo",
+      "power_law": "pl_hi",
+      "run_slug": "h_mid__d_lo__pl_hi",
+      "test_loss": 137.7809600830078,
+      "test_best_rerun_accuracy": null,
+      "test_best_rerun_mse": 144.83868408203125,
+      "test_triangles_total_structural": 1158.0,
+      "test_mse_by_total_triangles": 0.125076583835951,
+      "ood_test": {
+        "h_lo__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 134.12957763671875,
+          "test_triangles_total_structural": 1388.0,
+          "test_mse_by_total_triangles": 0.09663514238956682
+        },
+        "h_lo__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 31.069984436035156,
+          "test_triangles_total_structural": 190.0,
+          "test_mse_by_total_triangles": 0.16352623387386925
+        },
+        "h_lo__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 11623.2001953125,
+          "test_triangles_total_structural": 13185.0,
+          "test_mse_by_total_triangles": 0.8815472275549867
+        },
+        "h_lo__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 215.3048095703125,
+          "test_triangles_total_structural": 2967.0,
+          "test_mse_by_total_triangles": 0.07256650137186131
+        },
+        "h_mid__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 7578.6142578125,
+          "test_triangles_total_structural": 7485.0,
+          "test_mse_by_total_triangles": 1.0125069148714094
+        },
+        "h_mid__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 173480.078125,
+          "test_triangles_total_structural": 43064.0,
+          "test_mse_by_total_triangles": 4.0284246267183725
+        },
+        "h_mid__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 14914.7197265625,
+          "test_triangles_total_structural": 14262.0,
+          "test_mse_by_total_triangles": 1.0457663530053638
+        },
+        "h_hi__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 46251.01171875,
+          "test_triangles_total_structural": 19509.0,
+          "test_mse_by_total_triangles": 2.3707525613178535
+        },
+        "h_hi__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 3629.984130859375,
+          "test_triangles_total_structural": 5625.0,
+          "test_mse_by_total_triangles": 0.6453305121527778
+        },
+        "h_hi__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 757267.5625,
+          "test_triangles_total_structural": 88403.0,
+          "test_mse_by_total_triangles": 8.566084437179734
+        },
+        "h_hi__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 261409.859375,
+          "test_triangles_total_structural": 60093.0,
+          "test_mse_by_total_triangles": 4.350088352636746
+        }
+      },
+      "output_dir": "/raven/u/klettl/Projects/topobench/logs/train/runs/notebook_gu_grid_2026-07-29_18-09-01__triangle_counting__05__h_mid__d_lo__pl_hi__s43",
+      "wandb_config": {
+        "AvgTime/train_epoch_mean": 26.685278987884523,
+        "AvgTime/train_epoch_std": 0.015236519373931617,
+        "model/params/total": 680014,
+        "model/params/trainable": 680014,
+        "model/params/non_trainable": 0
+      }
+    },
+    {
+      "experiment": "triangle_counting",
+      "wandb_project": "challenge_triangle_counting",
+      "wandb_run_name": "gauge_hom_0.4-0.6__deg_1-2.5__gamma_4-5__s44",
+      "train_seed": 44,
+      "homophily": "h_mid",
+      "avg_degree": "d_lo",
+      "power_law": "pl_hi",
+      "run_slug": "h_mid__d_lo__pl_hi",
+      "test_loss": 136.29356384277344,
+      "test_best_rerun_accuracy": null,
+      "test_best_rerun_mse": 143.04754638671875,
+      "test_triangles_total_structural": 1158.0,
+      "test_mse_by_total_triangles": 0.12352983280372949,
+      "ood_test": {
+        "h_lo__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 125.86075592041016,
+          "test_triangles_total_structural": 1388.0,
+          "test_mse_by_total_triangles": 0.09067777804064132
+        },
+        "h_lo__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 50.16678237915039,
+          "test_triangles_total_structural": 190.0,
+          "test_mse_by_total_triangles": 0.2640356967323705
+        },
+        "h_lo__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 11303.7392578125,
+          "test_triangles_total_structural": 13185.0,
+          "test_mse_by_total_triangles": 0.8573181082906712
+        },
+        "h_lo__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 191.08538818359375,
+          "test_triangles_total_structural": 2967.0,
+          "test_mse_by_total_triangles": 0.06440356864967771
+        },
+        "h_mid__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 7383.45849609375,
+          "test_triangles_total_structural": 7485.0,
+          "test_mse_by_total_triangles": 0.9864340008141282
+        },
+        "h_mid__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 172121.453125,
+          "test_triangles_total_structural": 43064.0,
+          "test_mse_by_total_triangles": 3.996875653097715
+        },
+        "h_mid__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 14611.7587890625,
+          "test_triangles_total_structural": 14262.0,
+          "test_mse_by_total_triangles": 1.0245238247835156
+        },
+        "h_hi__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 45762.5234375,
+          "test_triangles_total_structural": 19509.0,
+          "test_mse_by_total_triangles": 2.3457134367471424
+        },
+        "h_hi__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 3500.013671875,
+          "test_triangles_total_structural": 5625.0,
+          "test_mse_by_total_triangles": 0.6222246527777778
+        },
+        "h_hi__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 754422.75,
+          "test_triangles_total_structural": 88403.0,
+          "test_mse_by_total_triangles": 8.533904392384875
+        },
+        "h_hi__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 259874.265625,
+          "test_triangles_total_structural": 60093.0,
+          "test_mse_by_total_triangles": 4.324534731582713
+        }
+      },
+      "output_dir": "/raven/u/klettl/Projects/topobench/logs/train/runs/notebook_gu_grid_2026-07-29_18-09-01__triangle_counting__05__h_mid__d_lo__pl_hi__s44",
+      "wandb_config": {
+        "AvgTime/train_epoch_mean": 26.746804816382273,
+        "AvgTime/train_epoch_std": 0.026101256131313384,
+        "model/params/total": 680014,
+        "model/params/trainable": 680014,
+        "model/params/non_trainable": 0
+      }
+    },
+    {
+      "experiment": "triangle_counting",
+      "wandb_project": "challenge_triangle_counting",
+      "wandb_run_name": "gauge_hom_0.4-0.6__deg_4-5__gamma_1.5-2__s42",
+      "train_seed": 42,
+      "homophily": "h_mid",
+      "avg_degree": "d_hi",
+      "power_law": "pl_lo",
+      "run_slug": "h_mid__d_hi__pl_lo",
+      "test_loss": 110186.390625,
+      "test_best_rerun_accuracy": null,
+      "test_best_rerun_mse": 77783.046875,
+      "test_triangles_total_structural": 43064.0,
+      "test_mse_by_total_triangles": 1.8062197398058704,
+      "ood_test": {
+        "h_lo__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 72692.9453125,
+          "test_triangles_total_structural": 1388.0,
+          "test_mse_by_total_triangles": 52.37243898595101
+        },
+        "h_lo__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 77523.5234375,
+          "test_triangles_total_structural": 190.0,
+          "test_mse_by_total_triangles": 408.01854440789475
+        },
+        "h_lo__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 41448.28515625,
+          "test_triangles_total_structural": 13185.0,
+          "test_mse_by_total_triangles": 3.143593868505878
+        },
+        "h_lo__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 71492.0703125,
+          "test_triangles_total_structural": 2967.0,
+          "test_mse_by_total_triangles": 24.095743280249412
+        },
+        "h_mid__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 55829.3359375,
+          "test_triangles_total_structural": 7485.0,
+          "test_mse_by_total_triangles": 7.458829116566466
+        },
+        "h_mid__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 73847.4921875,
+          "test_triangles_total_structural": 1158.0,
+          "test_mse_by_total_triangles": 63.771582199913645
+        },
+        "h_mid__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 45487.6640625,
+          "test_triangles_total_structural": 14262.0,
+          "test_mse_by_total_triangles": 3.18943093973496
+        },
+        "h_hi__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 45345.67578125,
+          "test_triangles_total_structural": 19509.0,
+          "test_mse_by_total_triangles": 2.3243464955277053
+        },
+        "h_hi__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 61831.02734375,
+          "test_triangles_total_structural": 5625.0,
+          "test_mse_by_total_triangles": 10.992182638888888
+        },
+        "h_hi__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 471568.59375,
+          "test_triangles_total_structural": 88403.0,
+          "test_mse_by_total_triangles": 5.334305326176714
+        },
+        "h_hi__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 122092.109375,
+          "test_triangles_total_structural": 60093.0,
+          "test_mse_by_total_triangles": 2.031719324630157
+        }
+      },
+      "output_dir": "/raven/u/klettl/Projects/topobench/logs/train/runs/notebook_gu_grid_2026-07-29_18-09-01__triangle_counting__06__h_mid__d_hi__pl_lo__s42",
+      "wandb_config": {
+        "AvgTime/train_epoch_mean": 26.76537295182546,
+        "AvgTime/train_epoch_std": 0.01980925732206768,
+        "model/params/total": 680014,
+        "model/params/trainable": 680014,
+        "model/params/non_trainable": 0
+      }
+    },
+    {
+      "experiment": "triangle_counting",
+      "wandb_project": "challenge_triangle_counting",
+      "wandb_run_name": "gauge_hom_0.4-0.6__deg_4-5__gamma_1.5-2__s43",
+      "train_seed": 43,
+      "homophily": "h_mid",
+      "avg_degree": "d_hi",
+      "power_law": "pl_lo",
+      "run_slug": "h_mid__d_hi__pl_lo",
+      "test_loss": 115386.234375,
+      "test_best_rerun_accuracy": null,
+      "test_best_rerun_mse": 80518.15625,
+      "test_triangles_total_structural": 43064.0,
+      "test_mse_by_total_triangles": 1.8697324040962289,
+      "ood_test": {
+        "h_lo__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 62194.53515625,
+          "test_triangles_total_structural": 1388.0,
+          "test_mse_by_total_triangles": 44.8087429079611
+        },
+        "h_lo__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 66429.046875,
+          "test_triangles_total_structural": 190.0,
+          "test_mse_by_total_triangles": 349.6265625
+        },
+        "h_lo__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 33251.91796875,
+          "test_triangles_total_structural": 13185.0,
+          "test_mse_by_total_triangles": 2.521950547497156
+        },
+        "h_lo__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 59275.30078125,
+          "test_triangles_total_structural": 2967.0,
+          "test_mse_by_total_triangles": 19.97819372472194
+        },
+        "h_mid__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 47307.64453125,
+          "test_triangles_total_structural": 7485.0,
+          "test_mse_by_total_triangles": 6.320326590681363
+        },
+        "h_mid__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 63277.33203125,
+          "test_triangles_total_structural": 1158.0,
+          "test_mse_by_total_triangles": 54.643637332685664
+        },
+        "h_mid__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 36974.59375,
+          "test_triangles_total_structural": 14262.0,
+          "test_mse_by_total_triangles": 2.592525154256065
+        },
+        "h_hi__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 40615.32421875,
+          "test_triangles_total_structural": 19509.0,
+          "test_mse_by_total_triangles": 2.0818762734507152
+        },
+        "h_hi__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 52273.48046875,
+          "test_triangles_total_structural": 5625.0,
+          "test_mse_by_total_triangles": 9.293063194444445
+        },
+        "h_hi__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 489638.875,
+          "test_triangles_total_structural": 88403.0,
+          "test_mse_by_total_triangles": 5.5387133355202876
+        },
+        "h_hi__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 127826.84375,
+          "test_triangles_total_structural": 60093.0,
+          "test_mse_by_total_triangles": 2.127150312848418
+        }
+      },
+      "output_dir": "/raven/u/klettl/Projects/topobench/logs/train/runs/notebook_gu_grid_2026-07-29_18-09-01__triangle_counting__06__h_mid__d_hi__pl_lo__s43",
+      "wandb_config": {
+        "AvgTime/train_epoch_mean": 30.1837486743927,
+        "AvgTime/train_epoch_std": 0.012592750999055522,
+        "model/params/total": 680014,
+        "model/params/trainable": 680014,
+        "model/params/non_trainable": 0
+      }
+    },
+    {
+      "experiment": "triangle_counting",
+      "wandb_project": "challenge_triangle_counting",
+      "wandb_run_name": "gauge_hom_0.4-0.6__deg_4-5__gamma_1.5-2__s44",
+      "train_seed": 44,
+      "homophily": "h_mid",
+      "avg_degree": "d_hi",
+      "power_law": "pl_lo",
+      "run_slug": "h_mid__d_hi__pl_lo",
+      "test_loss": 112996.703125,
+      "test_best_rerun_accuracy": null,
+      "test_best_rerun_mse": 79216.1171875,
+      "test_triangles_total_structural": 43064.0,
+      "test_mse_by_total_triangles": 1.8394974267950028,
+      "ood_test": {
+        "h_lo__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 67694.03125,
+          "test_triangles_total_structural": 1388.0,
+          "test_mse_by_total_triangles": 48.770915886167145
+        },
+        "h_lo__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 72372.703125,
+          "test_triangles_total_structural": 190.0,
+          "test_mse_by_total_triangles": 380.9089638157895
+        },
+        "h_lo__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 36571.75,
+          "test_triangles_total_structural": 13185.0,
+          "test_mse_by_total_triangles": 2.773739097459234
+        },
+        "h_lo__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 64882.30078125,
+          "test_triangles_total_structural": 2967.0,
+          "test_mse_by_total_triangles": 21.86798138902932
+        },
+        "h_mid__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 51803.01171875,
+          "test_triangles_total_structural": 7485.0,
+          "test_mse_by_total_triangles": 6.920910049265197
+        },
+        "h_mid__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 69307.5625,
+          "test_triangles_total_structural": 1158.0,
+          "test_mse_by_total_triangles": 59.8510902417962
+        },
+        "h_mid__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 40991.625,
+          "test_triangles_total_structural": 14262.0,
+          "test_mse_by_total_triangles": 2.874184896928902
+        },
+        "h_hi__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 43335.859375,
+          "test_triangles_total_structural": 19509.0,
+          "test_mse_by_total_triangles": 2.2213265351888873
+        },
+        "h_hi__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 57354.875,
+          "test_triangles_total_structural": 5625.0,
+          "test_mse_by_total_triangles": 10.196422222222223
+        },
+        "h_hi__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 478869.875,
+          "test_triangles_total_structural": 88403.0,
+          "test_mse_by_total_triangles": 5.416896202617559
+        },
+        "h_hi__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 124510.9453125,
+          "test_triangles_total_structural": 60093.0,
+          "test_mse_by_total_triangles": 2.071970867031102
+        }
+      },
+      "output_dir": "/raven/u/klettl/Projects/topobench/logs/train/runs/notebook_gu_grid_2026-07-29_18-09-01__triangle_counting__06__h_mid__d_hi__pl_lo__s44",
+      "wandb_config": {
+        "AvgTime/train_epoch_mean": 30.122872591018677,
+        "AvgTime/train_epoch_std": 0,
+        "model/params/total": 680014,
+        "model/params/trainable": 680014,
+        "model/params/non_trainable": 0
+      }
+    },
+    {
+      "experiment": "triangle_counting",
+      "wandb_project": "challenge_triangle_counting",
+      "wandb_run_name": "gauge_hom_0.4-0.6__deg_4-5__gamma_4-5__s42",
+      "train_seed": 42,
+      "homophily": "h_mid",
+      "avg_degree": "d_hi",
+      "power_law": "pl_hi",
+      "run_slug": "h_mid__d_hi__pl_hi",
+      "test_loss": 8051.67724609375,
+      "test_best_rerun_accuracy": null,
+      "test_best_rerun_mse": 7592.3466796875,
+      "test_triangles_total_structural": 14262.0,
+      "test_mse_by_total_triangles": 0.5323479652003575,
+      "ood_test": {
+        "h_lo__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 2092.030517578125,
+          "test_triangles_total_structural": 1388.0,
+          "test_mse_by_total_triangles": 1.5072265976787644
+        },
+        "h_lo__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 2085.065673828125,
+          "test_triangles_total_structural": 190.0,
+          "test_mse_by_total_triangles": 10.974029862253289
+        },
+        "h_lo__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 5833.92626953125,
+          "test_triangles_total_structural": 13185.0,
+          "test_mse_by_total_triangles": 0.44246691464021615
+        },
+        "h_lo__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 3703.797607421875,
+          "test_triangles_total_structural": 2967.0,
+          "test_mse_by_total_triangles": 1.2483308417330217
+        },
+        "h_mid__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 4426.53955078125,
+          "test_triangles_total_structural": 7485.0,
+          "test_mse_by_total_triangles": 0.5913880495365732
+        },
+        "h_mid__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 2031.8299560546875,
+          "test_triangles_total_structural": 1158.0,
+          "test_mse_by_total_triangles": 1.754602725435827
+        },
+        "h_mid__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 119591.828125,
+          "test_triangles_total_structural": 43064.0,
+          "test_mse_by_total_triangles": 2.7770719887841353
+        },
+        "h_hi__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 33156.43359375,
+          "test_triangles_total_structural": 19509.0,
+          "test_mse_by_total_triangles": 1.6995455222589575
+        },
+        "h_hi__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 3154.90478515625,
+          "test_triangles_total_structural": 5625.0,
+          "test_mse_by_total_triangles": 0.5608719618055555
+        },
+        "h_hi__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 646926.8125,
+          "test_triangles_total_structural": 88403.0,
+          "test_mse_by_total_triangles": 7.317928266009072
+        },
+        "h_hi__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 201174.671875,
+          "test_triangles_total_structural": 60093.0,
+          "test_mse_by_total_triangles": 3.34772222846255
+        }
+      },
+      "output_dir": "/raven/u/klettl/Projects/topobench/logs/train/runs/notebook_gu_grid_2026-07-29_18-09-01__triangle_counting__07__h_mid__d_hi__pl_hi__s42",
+      "wandb_config": {
+        "AvgTime/train_epoch_mean": 26.73908473054568,
+        "AvgTime/train_epoch_std": 0.017304739053562983,
+        "model/params/total": 680014,
+        "model/params/trainable": 680014,
+        "model/params/non_trainable": 0
+      }
+    },
+    {
+      "experiment": "triangle_counting",
+      "wandb_project": "challenge_triangle_counting",
+      "wandb_run_name": "gauge_hom_0.4-0.6__deg_4-5__gamma_4-5__s43",
+      "train_seed": 43,
+      "homophily": "h_mid",
+      "avg_degree": "d_hi",
+      "power_law": "pl_hi",
+      "run_slug": "h_mid__d_hi__pl_hi",
+      "test_loss": 8646.197265625,
+      "test_best_rerun_accuracy": null,
+      "test_best_rerun_mse": 8289.2724609375,
+      "test_triangles_total_structural": 14262.0,
+      "test_mse_by_total_triangles": 0.5812138873185738,
+      "ood_test": {
+        "h_lo__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 5410.5986328125,
+          "test_triangles_total_structural": 1388.0,
+          "test_mse_by_total_triangles": 3.8981258161473344
+        },
+        "h_lo__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 6247.08203125,
+          "test_triangles_total_structural": 190.0,
+          "test_mse_by_total_triangles": 32.879379111842105
+        },
+        "h_lo__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 4959.35791015625,
+          "test_triangles_total_structural": 13185.0,
+          "test_mse_by_total_triangles": 0.37613636026971936
+        },
+        "h_lo__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 4768.63037109375,
+          "test_triangles_total_structural": 2967.0,
+          "test_mse_by_total_triangles": 1.6072229090305865
+        },
+        "h_mid__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 6055.59521484375,
+          "test_triangles_total_structural": 7485.0,
+          "test_mse_by_total_triangles": 0.8090307568261523
+        },
+        "h_mid__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 5503.91943359375,
+          "test_triangles_total_structural": 1158.0,
+          "test_mse_by_total_triangles": 4.7529528787510795
+        },
+        "h_mid__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 131676.71875,
+          "test_triangles_total_structural": 43064.0,
+          "test_mse_by_total_triangles": 3.057698280466283
+        },
+        "h_hi__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 31523.56640625,
+          "test_triangles_total_structural": 19509.0,
+          "test_mse_by_total_triangles": 1.615847373327695
+        },
+        "h_hi__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 4647.56396484375,
+          "test_triangles_total_structural": 5625.0,
+          "test_mse_by_total_triangles": 0.82623359375
+        },
+        "h_hi__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 666528.125,
+          "test_triangles_total_structural": 88403.0,
+          "test_mse_by_total_triangles": 7.539655045643247
+        },
+        "h_hi__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 207047.875,
+          "test_triangles_total_structural": 60093.0,
+          "test_mse_by_total_triangles": 3.4454574576073753
+        }
+      },
+      "output_dir": "/raven/u/klettl/Projects/topobench/logs/train/runs/notebook_gu_grid_2026-07-29_18-09-01__triangle_counting__07__h_mid__d_hi__pl_hi__s43",
+      "wandb_config": {
+        "AvgTime/train_epoch_mean": 30.1066259543101,
+        "AvgTime/train_epoch_std": 0.019081035201831867,
+        "model/params/total": 680014,
+        "model/params/trainable": 680014,
+        "model/params/non_trainable": 0
+      }
+    },
+    {
+      "experiment": "triangle_counting",
+      "wandb_project": "challenge_triangle_counting",
+      "wandb_run_name": "gauge_hom_0.4-0.6__deg_4-5__gamma_4-5__s44",
+      "train_seed": 44,
+      "homophily": "h_mid",
+      "avg_degree": "d_hi",
+      "power_law": "pl_hi",
+      "run_slug": "h_mid__d_hi__pl_hi",
+      "test_loss": 8714.51171875,
+      "test_best_rerun_accuracy": null,
+      "test_best_rerun_mse": 8349.7646484375,
+      "test_triangles_total_structural": 14262.0,
+      "test_mse_by_total_triangles": 0.5854553813236222,
+      "ood_test": {
+        "h_lo__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 3125.315185546875,
+          "test_triangles_total_structural": 1388.0,
+          "test_mse_by_total_triangles": 2.251668001114463
+        },
+        "h_lo__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 3475.40185546875,
+          "test_triangles_total_structural": 190.0,
+          "test_mse_by_total_triangles": 18.29158871299342
+        },
+        "h_lo__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 5496.91552734375,
+          "test_triangles_total_structural": 13185.0,
+          "test_mse_by_total_triangles": 0.4169067521686576
+        },
+        "h_lo__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 3234.051513671875,
+          "test_triangles_total_structural": 2967.0,
+          "test_mse_by_total_triangles": 1.0900072509847911
+        },
+        "h_mid__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 5479.8818359375,
+          "test_triangles_total_structural": 7485.0,
+          "test_mse_by_total_triangles": 0.732115141741817
+        },
+        "h_mid__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 3121.449951171875,
+          "test_triangles_total_structural": 1158.0,
+          "test_mse_by_total_triangles": 2.69555263486345
+        },
+        "h_mid__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 135223.0625,
+          "test_triangles_total_structural": 43064.0,
+          "test_mse_by_total_triangles": 3.140048822682519
+        },
+        "h_hi__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 34573.5078125,
+          "test_triangles_total_structural": 19509.0,
+          "test_mse_by_total_triangles": 1.7721824702701316
+        },
+        "h_hi__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 3725.67822265625,
+          "test_triangles_total_structural": 5625.0,
+          "test_mse_by_total_triangles": 0.6623427951388889
+        },
+        "h_hi__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 679564.4375,
+          "test_triangles_total_structural": 88403.0,
+          "test_mse_by_total_triangles": 7.68711963960499
+        },
+        "h_hi__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 214580.9375,
+          "test_triangles_total_structural": 60093.0,
+          "test_mse_by_total_triangles": 3.570814196329023
+        }
+      },
+      "output_dir": "/raven/u/klettl/Projects/topobench/logs/train/runs/notebook_gu_grid_2026-07-29_18-09-01__triangle_counting__07__h_mid__d_hi__pl_hi__s44",
+      "wandb_config": {
+        "AvgTime/train_epoch_mean": 27.553436279296875,
+        "AvgTime/train_epoch_std": 1.3254523609456763,
+        "model/params/total": 680014,
+        "model/params/trainable": 680014,
+        "model/params/non_trainable": 0
+      }
+    },
+    {
+      "experiment": "triangle_counting",
+      "wandb_project": "challenge_triangle_counting",
+      "wandb_run_name": "gauge_hom_0.9-1__deg_1-2.5__gamma_1.5-2__s42",
+      "train_seed": 42,
+      "homophily": "h_hi",
+      "avg_degree": "d_lo",
+      "power_law": "pl_lo",
+      "run_slug": "h_hi__d_lo__pl_lo",
+      "test_loss": 16423.25,
+      "test_best_rerun_accuracy": null,
+      "test_best_rerun_mse": 17448.40625,
+      "test_triangles_total_structural": 19509.0,
+      "test_mse_by_total_triangles": 0.8943772745912143,
+      "ood_test": {
+        "h_lo__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 10492.66796875,
+          "test_triangles_total_structural": 1388.0,
+          "test_mse_by_total_triangles": 7.559559055295389
+        },
+        "h_lo__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 2594.142822265625,
+          "test_triangles_total_structural": 190.0,
+          "test_mse_by_total_triangles": 13.653383275082238
+        },
+        "h_lo__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 74047.8046875,
+          "test_triangles_total_structural": 13185.0,
+          "test_mse_by_total_triangles": 5.616064064277588
+        },
+        "h_lo__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 51558.15234375,
+          "test_triangles_total_structural": 2967.0,
+          "test_mse_by_total_triangles": 17.377199981041457
+        },
+        "h_mid__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 8031.88427734375,
+          "test_triangles_total_structural": 7485.0,
+          "test_mse_by_total_triangles": 1.073064031709252
+        },
+        "h_mid__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 4083.186279296875,
+          "test_triangles_total_structural": 1158.0,
+          "test_mse_by_total_triangles": 3.5260675987019647
+        },
+        "h_mid__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 48333.66796875,
+          "test_triangles_total_structural": 43064.0,
+          "test_mse_by_total_triangles": 1.122368288332482
+        },
+        "h_mid__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 36195.8359375,
+          "test_triangles_total_structural": 14262.0,
+          "test_mse_by_total_triangles": 2.537921465257327
+        },
+        "h_hi__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 5805.728515625,
+          "test_triangles_total_structural": 5625.0,
+          "test_mse_by_total_triangles": 1.0321295138888888
+        },
+        "h_hi__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 428989.34375,
+          "test_triangles_total_structural": 88403.0,
+          "test_mse_by_total_triangles": 4.852655947761954
+        },
+        "h_hi__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 102353.9765625,
+          "test_triangles_total_structural": 60093.0,
+          "test_mse_by_total_triangles": 1.7032595570615545
+        }
+      },
+      "output_dir": "/raven/u/klettl/Projects/topobench/logs/train/runs/notebook_gu_grid_2026-07-29_18-09-01__triangle_counting__08__h_hi__d_lo__pl_lo__s42",
+      "wandb_config": {
+        "AvgTime/train_epoch_mean": 26.623099054609025,
+        "AvgTime/train_epoch_std": 0.16658395067693374,
+        "model/params/total": 680014,
+        "model/params/trainable": 680014,
+        "model/params/non_trainable": 0
+      }
+    },
+    {
+      "experiment": "triangle_counting",
+      "wandb_project": "challenge_triangle_counting",
+      "wandb_run_name": "gauge_hom_0.9-1__deg_1-2.5__gamma_1.5-2__s43",
+      "train_seed": 43,
+      "homophily": "h_hi",
+      "avg_degree": "d_lo",
+      "power_law": "pl_lo",
+      "run_slug": "h_hi__d_lo__pl_lo",
+      "test_loss": 18641.939453125,
+      "test_best_rerun_accuracy": null,
+      "test_best_rerun_mse": 19125.494140625,
+      "test_triangles_total_structural": 19509.0,
+      "test_mse_by_total_triangles": 0.980342105726844,
+      "ood_test": {
+        "h_lo__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 9972.06640625,
+          "test_triangles_total_structural": 1388.0,
+          "test_mse_by_total_triangles": 7.184485883465418
+        },
+        "h_lo__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 3454.55517578125,
+          "test_triangles_total_structural": 190.0,
+          "test_mse_by_total_triangles": 18.181869346217105
+        },
+        "h_lo__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 85203.7890625,
+          "test_triangles_total_structural": 13185.0,
+          "test_mse_by_total_triangles": 6.4621758864239665
+        },
+        "h_lo__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 65276.4609375,
+          "test_triangles_total_structural": 2967.0,
+          "test_mse_by_total_triangles": 22.00082943629929
+        },
+        "h_mid__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 8448.94140625,
+          "test_triangles_total_structural": 7485.0,
+          "test_mse_by_total_triangles": 1.128783087007348
+        },
+        "h_mid__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 3741.719482421875,
+          "test_triangles_total_structural": 1158.0,
+          "test_mse_by_total_triangles": 3.2311912628859023
+        },
+        "h_mid__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 52513.0390625,
+          "test_triangles_total_structural": 43064.0,
+          "test_mse_by_total_triangles": 1.2194185180777448
+        },
+        "h_mid__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 45665.6953125,
+          "test_triangles_total_structural": 14262.0,
+          "test_mse_by_total_triangles": 3.2019138488641143
+        },
+        "h_hi__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 4850.65771484375,
+          "test_triangles_total_structural": 5625.0,
+          "test_mse_by_total_triangles": 0.8623391493055556
+        },
+        "h_hi__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 438620.96875,
+          "test_triangles_total_structural": 88403.0,
+          "test_mse_by_total_triangles": 4.961607284255059
+        },
+        "h_hi__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 107146.9140625,
+          "test_triangles_total_structural": 60093.0,
+          "test_mse_by_total_triangles": 1.7830182227963323
+        }
+      },
+      "output_dir": "/raven/u/klettl/Projects/topobench/logs/train/runs/notebook_gu_grid_2026-07-29_18-09-01__triangle_counting__08__h_hi__d_lo__pl_lo__s43",
+      "wandb_config": {
+        "AvgTime/train_epoch_mean": 26.520780324935913,
+        "AvgTime/train_epoch_std": 0.1651619948216022,
+        "model/params/total": 680014,
+        "model/params/trainable": 680014,
+        "model/params/non_trainable": 0
+      }
+    },
+    {
+      "experiment": "triangle_counting",
+      "wandb_project": "challenge_triangle_counting",
+      "wandb_run_name": "gauge_hom_0.9-1__deg_1-2.5__gamma_1.5-2__s44",
+      "train_seed": 44,
+      "homophily": "h_hi",
+      "avg_degree": "d_lo",
+      "power_law": "pl_lo",
+      "run_slug": "h_hi__d_lo__pl_lo",
+      "test_loss": 15914.2392578125,
+      "test_best_rerun_accuracy": null,
+      "test_best_rerun_mse": 16834.578125,
+      "test_triangles_total_structural": 19509.0,
+      "test_mse_by_total_triangles": 0.862913430980573,
+      "ood_test": {
+        "h_lo__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 9444.884765625,
+          "test_triangles_total_structural": 1388.0,
+          "test_mse_by_total_triangles": 6.80467202134366
+        },
+        "h_lo__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 4246.466796875,
+          "test_triangles_total_structural": 190.0,
+          "test_mse_by_total_triangles": 22.349825246710527
+        },
+        "h_lo__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 83767.34375,
+          "test_triangles_total_structural": 13185.0,
+          "test_mse_by_total_triangles": 6.353230470231323
+        },
+        "h_lo__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 60930.06640625,
+          "test_triangles_total_structural": 2967.0,
+          "test_mse_by_total_triangles": 20.53591722489046
+        },
+        "h_mid__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 9942.3056640625,
+          "test_triangles_total_structural": 7485.0,
+          "test_mse_by_total_triangles": 1.3282973499081496
+        },
+        "h_mid__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 6442.306640625,
+          "test_triangles_total_structural": 1158.0,
+          "test_mse_by_total_triangles": 5.563304525582901
+        },
+        "h_mid__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 46818.8671875,
+          "test_triangles_total_structural": 43064.0,
+          "test_mse_by_total_triangles": 1.0871927175250788
+        },
+        "h_mid__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 47150.84765625,
+          "test_triangles_total_structural": 14262.0,
+          "test_mse_by_total_triangles": 3.3060473745793018
+        },
+        "h_hi__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 7465.75,
+          "test_triangles_total_structural": 5625.0,
+          "test_mse_by_total_triangles": 1.3272444444444444
+        },
+        "h_hi__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 386600.84375,
+          "test_triangles_total_structural": 88403.0,
+          "test_mse_by_total_triangles": 4.373164301550853
+        },
+        "h_hi__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 96515.8046875,
+          "test_triangles_total_structural": 60093.0,
+          "test_mse_by_total_triangles": 1.6061072785099761
+        }
+      },
+      "output_dir": "/raven/u/klettl/Projects/topobench/logs/train/runs/notebook_gu_grid_2026-07-29_18-09-01__triangle_counting__08__h_hi__d_lo__pl_lo__s44",
+      "wandb_config": {
+        "AvgTime/train_epoch_mean": 26.57241090138753,
+        "AvgTime/train_epoch_std": 0.0464456355107282,
+        "model/params/total": 680014,
+        "model/params/trainable": 680014,
+        "model/params/non_trainable": 0
+      }
+    },
+    {
+      "experiment": "triangle_counting",
+      "wandb_project": "challenge_triangle_counting",
+      "wandb_run_name": "gauge_hom_0.9-1__deg_1-2.5__gamma_4-5__s42",
+      "train_seed": 42,
+      "homophily": "h_hi",
+      "avg_degree": "d_lo",
+      "power_law": "pl_hi",
+      "run_slug": "h_hi__d_lo__pl_hi",
+      "test_loss": 2379.95166015625,
+      "test_best_rerun_accuracy": null,
+      "test_best_rerun_mse": 2457.351806640625,
+      "test_triangles_total_structural": 5625.0,
+      "test_mse_by_total_triangles": 0.4368625434027778,
+      "ood_test": {
+        "h_lo__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 771.0359497070312,
+          "test_triangles_total_structural": 1388.0,
+          "test_mse_by_total_triangles": 0.5555014046880629
+        },
+        "h_lo__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 277.4176025390625,
+          "test_triangles_total_structural": 190.0,
+          "test_mse_by_total_triangles": 1.4600926449424343
+        },
+        "h_lo__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 3991.514892578125,
+          "test_triangles_total_structural": 13185.0,
+          "test_mse_by_total_triangles": 0.30273150493576984
+        },
+        "h_lo__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 2755.9794921875,
+          "test_triangles_total_structural": 2967.0,
+          "test_mse_by_total_triangles": 0.9288774830426356
+        },
+        "h_mid__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 4615.14697265625,
+          "test_triangles_total_structural": 7485.0,
+          "test_mse_by_total_triangles": 0.6165861018912826
+        },
+        "h_mid__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 375.8353576660156,
+          "test_triangles_total_structural": 1158.0,
+          "test_mse_by_total_triangles": 0.3245555765682346
+        },
+        "h_mid__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 114737.53125,
+          "test_triangles_total_structural": 43064.0,
+          "test_mse_by_total_triangles": 2.664349137330485
+        },
+        "h_mid__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 7533.4716796875,
+          "test_triangles_total_structural": 14262.0,
+          "test_mse_by_total_triangles": 0.5282198625499579
+        },
+        "h_hi__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 35923.18359375,
+          "test_triangles_total_structural": 19509.0,
+          "test_mse_by_total_triangles": 1.841364682646471
+        },
+        "h_hi__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 663999.5625,
+          "test_triangles_total_structural": 88403.0,
+          "test_mse_by_total_triangles": 7.511052368132304
+        },
+        "h_hi__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 209534.234375,
+          "test_triangles_total_structural": 60093.0,
+          "test_mse_by_total_triangles": 3.486832648977418
+        }
+      },
+      "output_dir": "/raven/u/klettl/Projects/topobench/logs/train/runs/notebook_gu_grid_2026-07-29_18-09-01__triangle_counting__09__h_hi__d_lo__pl_hi__s42",
+      "wandb_config": {
+        "AvgTime/train_epoch_mean": 26.35192687034607,
+        "AvgTime/train_epoch_std": 0.03360123548318638,
+        "model/params/total": 680014,
+        "model/params/trainable": 680014,
+        "model/params/non_trainable": 0
+      }
+    },
+    {
+      "experiment": "triangle_counting",
+      "wandb_project": "challenge_triangle_counting",
+      "wandb_run_name": "gauge_hom_0.9-1__deg_1-2.5__gamma_4-5__s43",
+      "train_seed": 43,
+      "homophily": "h_hi",
+      "avg_degree": "d_lo",
+      "power_law": "pl_hi",
+      "run_slug": "h_hi__d_lo__pl_hi",
+      "test_loss": 2557.224609375,
+      "test_best_rerun_accuracy": null,
+      "test_best_rerun_mse": 2650.076416015625,
+      "test_triangles_total_structural": 5625.0,
+      "test_mse_by_total_triangles": 0.4711246961805556,
+      "ood_test": {
+        "h_lo__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 779.063232421875,
+          "test_triangles_total_structural": 1388.0,
+          "test_mse_by_total_triangles": 0.5612847495834834
+        },
+        "h_lo__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 378.194091796875,
+          "test_triangles_total_structural": 190.0,
+          "test_mse_by_total_triangles": 1.9904952199835526
+        },
+        "h_lo__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 4216.45166015625,
+          "test_triangles_total_structural": 13185.0,
+          "test_mse_by_total_triangles": 0.3197915555674061
+        },
+        "h_lo__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 2763.764892578125,
+          "test_triangles_total_structural": 2967.0,
+          "test_mse_by_total_triangles": 0.931501480477966
+        },
+        "h_mid__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 5185.3583984375,
+          "test_triangles_total_structural": 7485.0,
+          "test_mse_by_total_triangles": 0.6927666530978623
+        },
+        "h_mid__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 452.92791748046875,
+          "test_triangles_total_structural": 1158.0,
+          "test_mse_by_total_triangles": 0.39112946241836677
+        },
+        "h_mid__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 130458.796875,
+          "test_triangles_total_structural": 43064.0,
+          "test_mse_by_total_triangles": 3.0294166095810886
+        },
+        "h_mid__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 8104.529296875,
+          "test_triangles_total_structural": 14262.0,
+          "test_mse_by_total_triangles": 0.5682603629838031
+        },
+        "h_hi__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 39056.97265625,
+          "test_triangles_total_structural": 19509.0,
+          "test_mse_by_total_triangles": 2.001997675752217
+        },
+        "h_hi__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 687038.8125,
+          "test_triangles_total_structural": 88403.0,
+          "test_mse_by_total_triangles": 7.771668523692635
+        },
+        "h_hi__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 220042.75,
+          "test_triangles_total_structural": 60093.0,
+          "test_mse_by_total_triangles": 3.661703526201055
+        }
+      },
+      "output_dir": "/raven/u/klettl/Projects/topobench/logs/train/runs/notebook_gu_grid_2026-07-29_18-09-01__triangle_counting__09__h_hi__d_lo__pl_hi__s43",
+      "wandb_config": {
+        "AvgTime/train_epoch_mean": 26.345271229743958,
+        "AvgTime/train_epoch_std": 0.023256148807822948,
+        "model/params/total": 680014,
+        "model/params/trainable": 680014,
+        "model/params/non_trainable": 0
+      }
+    },
+    {
+      "experiment": "triangle_counting",
+      "wandb_project": "challenge_triangle_counting",
+      "wandb_run_name": "gauge_hom_0.9-1__deg_1-2.5__gamma_4-5__s44",
+      "train_seed": 44,
+      "homophily": "h_hi",
+      "avg_degree": "d_lo",
+      "power_law": "pl_hi",
+      "run_slug": "h_hi__d_lo__pl_hi",
+      "test_loss": 2620.078369140625,
+      "test_best_rerun_accuracy": null,
+      "test_best_rerun_mse": 2776.860107421875,
+      "test_triangles_total_structural": 5625.0,
+      "test_mse_by_total_triangles": 0.49366401909722224,
+      "ood_test": {
+        "h_lo__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 562.0873413085938,
+          "test_triangles_total_structural": 1388.0,
+          "test_mse_by_total_triangles": 0.40496206146152286
+        },
+        "h_lo__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 842.1431274414062,
+          "test_triangles_total_structural": 190.0,
+          "test_mse_by_total_triangles": 4.432332249691612
+        },
+        "h_lo__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 8354.6025390625,
+          "test_triangles_total_structural": 13185.0,
+          "test_mse_by_total_triangles": 0.6336444853289723
+        },
+        "h_lo__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 328.3114013671875,
+          "test_triangles_total_structural": 2967.0,
+          "test_mse_by_total_triangles": 0.1106543314348458
+        },
+        "h_mid__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 5917.68603515625,
+          "test_triangles_total_structural": 7485.0,
+          "test_mse_by_total_triangles": 0.7906060167209419
+        },
+        "h_mid__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 698.205810546875,
+          "test_triangles_total_structural": 1158.0,
+          "test_mse_by_total_triangles": 0.6029411144618955
+        },
+        "h_mid__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 159688.453125,
+          "test_triangles_total_structural": 43064.0,
+          "test_mse_by_total_triangles": 3.7081658258638304
+        },
+        "h_mid__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 11597.923828125,
+          "test_triangles_total_structural": 14262.0,
+          "test_mse_by_total_triangles": 0.8132045875841397
+        },
+        "h_hi__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 40327.828125,
+          "test_triangles_total_structural": 19509.0,
+          "test_mse_by_total_triangles": 2.0671396855297557
+        },
+        "h_hi__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 726290.75,
+          "test_triangles_total_structural": 88403.0,
+          "test_mse_by_total_triangles": 8.215679897741026
+        },
+        "h_hi__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 243534.859375,
+          "test_triangles_total_structural": 60093.0,
+          "test_mse_by_total_triangles": 4.052632742166309
+        }
+      },
+      "output_dir": "/raven/u/klettl/Projects/topobench/logs/train/runs/notebook_gu_grid_2026-07-29_18-09-01__triangle_counting__09__h_hi__d_lo__pl_hi__s44",
+      "wandb_config": {
+        "AvgTime/train_epoch_mean": 26.551474452018738,
+        "AvgTime/train_epoch_std": 0.00813901424407959,
+        "model/params/total": 680014,
+        "model/params/trainable": 680014,
+        "model/params/non_trainable": 0
+      }
+    },
+    {
+      "experiment": "triangle_counting",
+      "wandb_project": "challenge_triangle_counting",
+      "wandb_run_name": "gauge_hom_0.9-1__deg_4-5__gamma_1.5-2__s42",
+      "train_seed": 42,
+      "homophily": "h_hi",
+      "avg_degree": "d_hi",
+      "power_law": "pl_lo",
+      "run_slug": "h_hi__d_hi__pl_lo",
+      "test_loss": 504146.5625,
+      "test_best_rerun_accuracy": null,
+      "test_best_rerun_mse": 325134.40625,
+      "test_triangles_total_structural": 88403.0,
+      "test_mse_by_total_triangles": 3.6778662064635816,
+      "ood_test": {
+        "h_lo__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 375161.4375,
+          "test_triangles_total_structural": 1388.0,
+          "test_mse_by_total_triangles": 270.28922010086455
+        },
+        "h_lo__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 387830.5,
+          "test_triangles_total_structural": 190.0,
+          "test_mse_by_total_triangles": 2041.213157894737
+        },
+        "h_lo__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 304404.46875,
+          "test_triangles_total_structural": 13185.0,
+          "test_mse_by_total_triangles": 23.08718003412969
+        },
+        "h_lo__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 394760.09375,
+          "test_triangles_total_structural": 2967.0,
+          "test_mse_by_total_triangles": 133.05025067408155
+        },
+        "h_mid__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 327185.6875,
+          "test_triangles_total_structural": 7485.0,
+          "test_mse_by_total_triangles": 43.71218269873079
+        },
+        "h_mid__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 378895.625,
+          "test_triangles_total_structural": 1158.0,
+          "test_mse_by_total_triangles": 327.1982944732297
+        },
+        "h_mid__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 176803.75,
+          "test_triangles_total_structural": 43064.0,
+          "test_mse_by_total_triangles": 4.1056044491919
+        },
+        "h_mid__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 313371.78125,
+          "test_triangles_total_structural": 14262.0,
+          "test_mse_by_total_triangles": 21.97249903589959
+        },
+        "h_hi__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 254339.625,
+          "test_triangles_total_structural": 19509.0,
+          "test_mse_by_total_triangles": 13.037040596647701
+        },
+        "h_hi__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 348246.03125,
+          "test_triangles_total_structural": 5625.0,
+          "test_mse_by_total_triangles": 61.910405555555556
+        },
+        "h_hi__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 162495.5625,
+          "test_triangles_total_structural": 60093.0,
+          "test_mse_by_total_triangles": 2.704068069492287
+        }
+      },
+      "output_dir": "/raven/u/klettl/Projects/topobench/logs/train/runs/notebook_gu_grid_2026-07-29_18-09-01__triangle_counting__10__h_hi__d_hi__pl_lo__s42",
+      "wandb_config": {
+        "AvgTime/train_epoch_mean": 26.7362412661314,
+        "AvgTime/train_epoch_std": 0.022386848377965236,
+        "model/params/total": 680014,
+        "model/params/trainable": 680014,
+        "model/params/non_trainable": 0
+      }
+    },
+    {
+      "experiment": "triangle_counting",
+      "wandb_project": "challenge_triangle_counting",
+      "wandb_run_name": "gauge_hom_0.9-1__deg_4-5__gamma_1.5-2__s43",
+      "train_seed": 43,
+      "homophily": "h_hi",
+      "avg_degree": "d_hi",
+      "power_law": "pl_lo",
+      "run_slug": "h_hi__d_hi__pl_lo",
+      "test_loss": 517046.59375,
+      "test_best_rerun_accuracy": null,
+      "test_best_rerun_mse": 329081.375,
+      "test_triangles_total_structural": 88403.0,
+      "test_mse_by_total_triangles": 3.7225136590387202,
+      "ood_test": {
+        "h_lo__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 320284.875,
+          "test_triangles_total_structural": 1388.0,
+          "test_mse_by_total_triangles": 230.75279178674353
+        },
+        "h_lo__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 323652.90625,
+          "test_triangles_total_structural": 190.0,
+          "test_mse_by_total_triangles": 1703.4363486842105
+        },
+        "h_lo__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 286870.4375,
+          "test_triangles_total_structural": 13185.0,
+          "test_mse_by_total_triangles": 21.75733314372393
+        },
+        "h_lo__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 371696.78125,
+          "test_triangles_total_structural": 2967.0,
+          "test_mse_by_total_triangles": 125.2769737950792
+        },
+        "h_mid__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 277863.1875,
+          "test_triangles_total_structural": 7485.0,
+          "test_mse_by_total_triangles": 37.12267034068136
+        },
+        "h_mid__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 318693.875,
+          "test_triangles_total_structural": 1158.0,
+          "test_mse_by_total_triangles": 275.21060017271157
+        },
+        "h_mid__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 160311.984375,
+          "test_triangles_total_structural": 43064.0,
+          "test_mse_by_total_triangles": 3.722645002206019
+        },
+        "h_mid__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 291439.1875,
+          "test_triangles_total_structural": 14262.0,
+          "test_mse_by_total_triangles": 20.43466466834946
+        },
+        "h_hi__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 210292.6875,
+          "test_triangles_total_structural": 19509.0,
+          "test_mse_by_total_triangles": 10.779265339074273
+        },
+        "h_hi__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 298111.625,
+          "test_triangles_total_structural": 5625.0,
+          "test_mse_by_total_triangles": 52.99762222222222
+        },
+        "h_hi__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 150024.859375,
+          "test_triangles_total_structural": 60093.0,
+          "test_mse_by_total_triangles": 2.4965446786647365
+        }
+      },
+      "output_dir": "/raven/u/klettl/Projects/topobench/logs/train/runs/notebook_gu_grid_2026-07-29_18-09-01__triangle_counting__10__h_hi__d_hi__pl_lo__s43",
+      "wandb_config": {
+        "AvgTime/train_epoch_mean": 27.0877879517419,
+        "AvgTime/train_epoch_std": 0.01967982288814584,
+        "model/params/total": 680014,
+        "model/params/trainable": 680014,
+        "model/params/non_trainable": 0
+      }
+    },
+    {
+      "experiment": "triangle_counting",
+      "wandb_project": "challenge_triangle_counting",
+      "wandb_run_name": "gauge_hom_0.9-1__deg_4-5__gamma_1.5-2__s44",
+      "train_seed": 44,
+      "homophily": "h_hi",
+      "avg_degree": "d_hi",
+      "power_law": "pl_lo",
+      "run_slug": "h_hi__d_hi__pl_lo",
+      "test_loss": 520948.8125,
+      "test_best_rerun_accuracy": null,
+      "test_best_rerun_mse": 331958.53125,
+      "test_triangles_total_structural": 88403.0,
+      "test_mse_by_total_triangles": 3.7550595709421626,
+      "ood_test": {
+        "h_lo__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 348509.96875,
+          "test_triangles_total_structural": 1388.0,
+          "test_mse_by_total_triangles": 251.08787373919307
+        },
+        "h_lo__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 353351.25,
+          "test_triangles_total_structural": 190.0,
+          "test_mse_by_total_triangles": 1859.7434210526317
+        },
+        "h_lo__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 304730.96875,
+          "test_triangles_total_structural": 13185.0,
+          "test_mse_by_total_triangles": 23.11194302237391
+        },
+        "h_lo__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 388165.65625,
+          "test_triangles_total_structural": 2967.0,
+          "test_mse_by_total_triangles": 130.82765630266263
+        },
+        "h_mid__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 303372.96875,
+          "test_triangles_total_structural": 7485.0,
+          "test_mse_by_total_triangles": 40.53079074816299
+        },
+        "h_mid__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 348204.15625,
+          "test_triangles_total_structural": 1158.0,
+          "test_mse_by_total_triangles": 300.6944354490501
+        },
+        "h_mid__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 170844.6875,
+          "test_triangles_total_structural": 43064.0,
+          "test_mse_by_total_triangles": 3.967227556659855
+        },
+        "h_mid__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 308185.90625,
+          "test_triangles_total_structural": 14262.0,
+          "test_mse_by_total_triangles": 21.60888418524751
+        },
+        "h_hi__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 232932.34375,
+          "test_triangles_total_structural": 19509.0,
+          "test_mse_by_total_triangles": 11.939737749243939
+        },
+        "h_hi__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 325812.25,
+          "test_triangles_total_structural": 5625.0,
+          "test_mse_by_total_triangles": 57.922177777777776
+        },
+        "h_hi__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 156855.59375,
+          "test_triangles_total_structural": 60093.0,
+          "test_mse_by_total_triangles": 2.6102140640340803
+        }
+      },
+      "output_dir": "/raven/u/klettl/Projects/topobench/logs/train/runs/notebook_gu_grid_2026-07-29_18-09-01__triangle_counting__10__h_hi__d_hi__pl_lo__s44",
+      "wandb_config": {
+        "AvgTime/train_epoch_mean": 30.104320398966472,
+        "AvgTime/train_epoch_std": 0.8285811505890811,
+        "model/params/total": 680014,
+        "model/params/trainable": 680014,
+        "model/params/non_trainable": 0
+      }
+    },
+    {
+      "experiment": "triangle_counting",
+      "wandb_project": "challenge_triangle_counting",
+      "wandb_run_name": "gauge_hom_0.9-1__deg_4-5__gamma_4-5__s42",
+      "train_seed": 42,
+      "homophily": "h_hi",
+      "avg_degree": "d_hi",
+      "power_law": "pl_hi",
+      "run_slug": "h_hi__d_hi__pl_hi",
+      "test_loss": 112100.7109375,
+      "test_best_rerun_accuracy": null,
+      "test_best_rerun_mse": 107883.7109375,
+      "test_triangles_total_structural": 60093.0,
+      "test_mse_by_total_triangles": 1.7952791662506449,
+      "ood_test": {
+        "h_lo__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 144999.140625,
+          "test_triangles_total_structural": 1388.0,
+          "test_mse_by_total_triangles": 104.46623964337176
+        },
+        "h_lo__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 151683.875,
+          "test_triangles_total_structural": 190.0,
+          "test_mse_by_total_triangles": 798.3361842105263
+        },
+        "h_lo__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 98781.9921875,
+          "test_triangles_total_structural": 13185.0,
+          "test_mse_by_total_triangles": 7.491997890595374
+        },
+        "h_lo__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 146281.890625,
+          "test_triangles_total_structural": 2967.0,
+          "test_mse_by_total_triangles": 49.30296279912369
+        },
+        "h_mid__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 116696.015625,
+          "test_triangles_total_structural": 7485.0,
+          "test_mse_by_total_triangles": 15.5906500501002
+        },
+        "h_mid__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 145396.234375,
+          "test_triangles_total_structural": 1158.0,
+          "test_mse_by_total_triangles": 125.55806077288429
+        },
+        "h_mid__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 82086.046875,
+          "test_triangles_total_structural": 43064.0,
+          "test_mse_by_total_triangles": 1.9061407875487646
+        },
+        "h_mid__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 102296.1015625,
+          "test_triangles_total_structural": 14262.0,
+          "test_mse_by_total_triangles": 7.172633681285935
+        },
+        "h_hi__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 86106.0078125,
+          "test_triangles_total_structural": 19509.0,
+          "test_mse_by_total_triangles": 4.413655636501102
+        },
+        "h_hi__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 126764.6171875,
+          "test_triangles_total_structural": 5625.0,
+          "test_mse_by_total_triangles": 22.535931944444446
+        },
+        "h_hi__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 400886.625,
+          "test_triangles_total_structural": 88403.0,
+          "test_mse_by_total_triangles": 4.534762677737181
+        }
+      },
+      "output_dir": "/raven/u/klettl/Projects/topobench/logs/train/runs/notebook_gu_grid_2026-07-29_18-09-01__triangle_counting__11__h_hi__d_hi__pl_hi__s42",
+      "wandb_config": {
+        "AvgTime/train_epoch_mean": 27.021645409720286,
+        "AvgTime/train_epoch_std": 0.015819469744172632,
+        "model/params/total": 680014,
+        "model/params/trainable": 680014,
+        "model/params/non_trainable": 0
+      }
+    },
+    {
+      "experiment": "triangle_counting",
+      "wandb_project": "challenge_triangle_counting",
+      "wandb_run_name": "gauge_hom_0.9-1__deg_4-5__gamma_4-5__s43",
+      "train_seed": 43,
+      "homophily": "h_hi",
+      "avg_degree": "d_hi",
+      "power_law": "pl_hi",
+      "run_slug": "h_hi__d_hi__pl_hi",
+      "test_loss": 110922.265625,
+      "test_best_rerun_accuracy": null,
+      "test_best_rerun_mse": 105510.4140625,
+      "test_triangles_total_structural": 60093.0,
+      "test_mse_by_total_triangles": 1.7557854336195564,
+      "ood_test": {
+        "h_lo__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 130699.2265625,
+          "test_triangles_total_structural": 1388.0,
+          "test_mse_by_total_triangles": 94.16370789805475
+        },
+        "h_lo__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 136233.796875,
+          "test_triangles_total_structural": 190.0,
+          "test_mse_by_total_triangles": 717.0199835526316
+        },
+        "h_lo__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 93475.578125,
+          "test_triangles_total_structural": 13185.0,
+          "test_mse_by_total_triangles": 7.089539486158514
+        },
+        "h_lo__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 139833.96875,
+          "test_triangles_total_structural": 2967.0,
+          "test_mse_by_total_triangles": 47.12975016852039
+        },
+        "h_mid__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 104439.1328125,
+          "test_triangles_total_structural": 7485.0,
+          "test_mse_by_total_triangles": 13.953123956245825
+        },
+        "h_mid__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 132155.4375,
+          "test_triangles_total_structural": 1158.0,
+          "test_mse_by_total_triangles": 114.12386658031087
+        },
+        "h_mid__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 79156.3828125,
+          "test_triangles_total_structural": 43064.0,
+          "test_mse_by_total_triangles": 1.8381103198146944
+        },
+        "h_mid__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 96885.1171875,
+          "test_triangles_total_structural": 14262.0,
+          "test_mse_by_total_triangles": 6.793234973180479
+        },
+        "h_hi__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 78476.296875,
+          "test_triangles_total_structural": 19509.0,
+          "test_mse_by_total_triangles": 4.022568910502845
+        },
+        "h_hi__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 117785.9765625,
+          "test_triangles_total_structural": 5625.0,
+          "test_mse_by_total_triangles": 20.939729166666666
+        },
+        "h_hi__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 409772.75,
+          "test_triangles_total_structural": 88403.0,
+          "test_mse_by_total_triangles": 4.6352810424985575
+        }
+      },
+      "output_dir": "/raven/u/klettl/Projects/topobench/logs/train/runs/notebook_gu_grid_2026-07-29_18-09-01__triangle_counting__11__h_hi__d_hi__pl_hi__s43",
+      "wandb_config": {
+        "AvgTime/train_epoch_mean": 27.27666376431783,
+        "AvgTime/train_epoch_std": 0.4140196140369854,
+        "model/params/total": 680014,
+        "model/params/trainable": 680014,
+        "model/params/non_trainable": 0
+      }
+    },
+    {
+      "experiment": "triangle_counting",
+      "wandb_project": "challenge_triangle_counting",
+      "wandb_run_name": "gauge_hom_0.9-1__deg_4-5__gamma_4-5__s44",
+      "train_seed": 44,
+      "homophily": "h_hi",
+      "avg_degree": "d_hi",
+      "power_law": "pl_hi",
+      "run_slug": "h_hi__d_hi__pl_hi",
+      "test_loss": 111957.1484375,
+      "test_best_rerun_accuracy": null,
+      "test_best_rerun_mse": 108015.703125,
+      "test_triangles_total_structural": 60093.0,
+      "test_mse_by_total_triangles": 1.7974756315211422,
+      "ood_test": {
+        "h_lo__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 134526.484375,
+          "test_triangles_total_structural": 1388.0,
+          "test_mse_by_total_triangles": 96.92109825288185
+        },
+        "h_lo__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 141746.265625,
+          "test_triangles_total_structural": 190.0,
+          "test_mse_by_total_triangles": 746.0329769736842
+        },
+        "h_lo__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 89033.7265625,
+          "test_triangles_total_structural": 13185.0,
+          "test_mse_by_total_triangles": 6.752652754076602
+        },
+        "h_lo__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 135758.890625,
+          "test_triangles_total_structural": 2967.0,
+          "test_mse_by_total_triangles": 45.75628265082575
+        },
+        "h_mid__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 108287.484375,
+          "test_triangles_total_structural": 7485.0,
+          "test_mse_by_total_triangles": 14.467265781563126
+        },
+        "h_mid__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 136000.59375,
+          "test_triangles_total_structural": 1158.0,
+          "test_mse_by_total_triangles": 117.44438147668394
+        },
+        "h_mid__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 79795.40625,
+          "test_triangles_total_structural": 43064.0,
+          "test_mse_by_total_triangles": 1.8529492441482445
+        },
+        "h_mid__d_hi__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 93660.234375,
+          "test_triangles_total_structural": 14262.0,
+          "test_mse_by_total_triangles": 6.567117821834245
+        },
+        "h_hi__d_lo__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 80316.1171875,
+          "test_triangles_total_structural": 19509.0,
+          "test_mse_by_total_triangles": 4.116875144164232
+        },
+        "h_hi__d_lo__pl_hi": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 117668.1328125,
+          "test_triangles_total_structural": 5625.0,
+          "test_mse_by_total_triangles": 20.918779166666667
+        },
+        "h_hi__d_hi__pl_lo": {
+          "test_best_rerun_accuracy": null,
+          "test_best_rerun_mse": 406826.09375,
+          "test_triangles_total_structural": 88403.0,
+          "test_mse_by_total_triangles": 4.601948958180152
+        }
+      },
+      "output_dir": "/raven/u/klettl/Projects/topobench/logs/train/runs/notebook_gu_grid_2026-07-29_18-09-01__triangle_counting__11__h_hi__d_hi__pl_hi__s44",
+      "wandb_config": {
+        "AvgTime/train_epoch_mean": 30.593818366527557,
+        "AvgTime/train_epoch_std": 0.144298228729869,
+        "model/params/total": 680014,
+        "model/params/trainable": 680014,
+        "model/params/non_trainable": 0
+      }
+    }
+  ]
+}
diff --git a/2026_tdl_challenge/run_evaluation.ipynb b/2026_tdl_challenge/run_evaluation.ipynb
index 8542dbaab..ad6c01ef6 100644
--- a/2026_tdl_challenge/run_evaluation.ipynb
+++ b/2026_tdl_challenge/run_evaluation.ipynb
@@ -98,13 +98,13 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": null,
    "id": "config_cell",
    "metadata": {},
    "outputs": [],
    "source": [
     "# Your model configuration (e.g., \"graph/gcn\", \"graph/gin\", \"graph/gat\")\n",
-    "MODEL_CONFIG = \"graph/gin\""
+    "MODEL_CONFIG = \"graph/gauge\""
    ]
   },
   {
diff --git a/configs/model/graph/gauge.yaml b/configs/model/graph/gauge.yaml
new file mode 100644
index 000000000..8c267e301
--- /dev/null
+++ b/configs/model/graph/gauge.yaml
@@ -0,0 +1,52 @@
+_target_: topobench.model.TBModel
+
+model_name: gauge
+model_domain: graph
+
+feature_encoder:
+  _target_: topobench.nn.encoders.${model.feature_encoder.encoder_name}
+  encoder_name: AllCellFeatureEncoder
+  in_channels: ${infer_in_channels:${dataset},${oc.select:transforms,null}}
+  out_channels: 64
+  proj_dropout: 0.0
+
+backbone:
+  _target_: topobench.nn.backbones.GaugeModel
+  n_layers: 2
+  in_channels: ${model.feature_encoder.out_channels}
+  d_embedd: 128
+  r: 16                                               # frame dim; must satisfy r <= d_embedd
+  n_gated: 2
+  gamma: 0.01
+  tau: 1.0
+  bias: true
+  act: gelu                                          # feed-forward activation: relu|leaky_relu|gelu|sigmoid
+  f_sim_act: leaky_relu                              # similarity-scorer (f_sim) activation: relu|leaky_relu|gelu|sigmoid
+  dropout: 0.1                                        # dropout in the feed-forward blocks (fflayer + phi); reference default
+  f_sim_dropout: 0.0                                 # dropout in the similarity-scorer (f_sim)
+  phi_hidden_layers: 0                               # 0 -> single linear residual; null -> disable residual (reference behavior)
+  phi_hidden_dim: null                               # null -> defaults to d_embedd (unused when phi_hidden_layers == 0)
+  loss:                                              # Dirichlet-energy regularizer, added to the task loss (see configs/loss/default.yaml)
+    _target_: topobench.loss.model.DirichletLoss
+    lamb: 0.1                                        # lambda; weight of the Dirichlet-energy term
+    reduction: mean                                  # neighbor aggregation: mean|sum
+
+backbone_wrapper:
+  _target_: topobench.nn.wrappers.GaugeWrapper
+  _partial_: true
+  wrapper_name: GaugeWrapper
+  out_channels: ${model.backbone.d_embedd}
+  residual_connections: false   # not part of the paper; also lets d_embedd differ from the encoder width
+  num_cell_dimensions: ${infer_num_cell_dimensions:${oc.select:model.feature_encoder.selected_dimensions,null},${model.feature_encoder.in_channels}}
+
+readout:
+  _target_: topobench.nn.readouts.${model.readout.readout_name}
+  readout_name: NoReadOut #  Use <NoReadOut> in case readout is not needed Options: PropagateSignalDown
+  num_cell_dimensions: ${infer_num_cell_dimensions:${oc.select:model.feature_encoder.selected_dimensions,null},${model.feature_encoder.in_channels}} # The highest order of cell dimensions to consider
+  hidden_dim: ${model.backbone.d_embedd}   # readout consumes the backbone output width
+  out_channels: ${dataset.parameters.num_classes}
+  task_level: ${define_task_level:${dataset.parameters.task_level},${dataset.split_params.learning_setting}} # Handles the edge case of node-inductive task
+  pooling_type: sum
+
+# compile model for faster training with pytorch 2.0
+compile: false
diff --git a/test/loss/test_dirichlet_loss.py b/test/loss/test_dirichlet_loss.py
new file mode 100644
index 000000000..cd548e19f
--- /dev/null
+++ b/test/loss/test_dirichlet_loss.py
@@ -0,0 +1,253 @@
+"""Test the DirichletLoss class."""
+
+import pytest
+import torch
+import torch_geometric
+
+# Import from the package (populated by the loss-discovery mechanism), not the
+# submodule: `from ...model.DirichletLoss import DirichletLoss` would force-import
+# the submodule and shadow the class registered on the package, breaking hydra's
+# `_target_: topobench.loss.model.DirichletLoss` resolution in later tests.
+from topobench.loss.model import DirichletLoss
+
+
+def _make_inputs(N=3, r=2, d=4, edge_index=None, requires_grad=False):
+    """Build a ``(model_out, batch)`` pair for the Dirichlet loss.
+
+    Parameters
+    ----------
+    N : int, optional
+        Number of nodes (default: 3).
+    r : int, optional
+        Number of frame vectors (default: 2).
+    d : int, optional
+        Embedding dimension (default: 4).
+    edge_index : torch.Tensor or None, optional
+        Edge index of shape ``[2, E]``. Defaults to a directed cycle over the
+        ``N`` nodes so that every node has exactly one incoming edge.
+    requires_grad : bool, optional
+        Whether the embedding/frame tensors require gradients (default: False).
+
+    Returns
+    -------
+    tuple of (dict, torch_geometric.data.Data)
+        The mock model output and batch.
+    """
+    if edge_index is None:
+        src = torch.arange(N)
+        dst = torch.roll(src, -1)
+        edge_index = torch.stack([src, dst], dim=0)
+
+    model_out = {
+        "x_0": torch.randn(N, d, requires_grad=requires_grad),
+        "z_0": torch.randn(N, d, requires_grad=requires_grad),
+        "Q": torch.randn(N, r, d, requires_grad=requires_grad),
+    }
+    batch = torch_geometric.data.Data(edge_index=edge_index, num_nodes=N)
+    return model_out, batch
+
+
+def test_dirichlet_loss_init():
+    """Default hyperparameters are stored as given."""
+    loss_fn = DirichletLoss()
+    assert loss_fn.lamb == 0.1
+    assert loss_fn.reduce == "mean"
+
+
+def test_dirichlet_loss_init_invalid_reduction():
+    """An unsupported reduction raises ``NotImplementedError``."""
+    with pytest.raises(NotImplementedError):
+        DirichletLoss(reduction="max")
+
+
+def test_dirichlet_loss_repr():
+    """The repr reports the configured hyperparameters."""
+    assert repr(DirichletLoss()) == "DirichletLoss(lamb=0.1, reduction=mean)"
+
+
+def test_dirichlet_loss_forward_is_nonnegative_scalar():
+    """The forward pass returns a non-negative scalar tensor."""
+    loss_fn = DirichletLoss()
+    model_out, batch = _make_inputs()
+    loss = loss_fn.forward(model_out, batch)
+    assert isinstance(loss, torch.Tensor)
+    assert loss.dim() == 0
+    assert loss.item() >= 0.0
+
+
+@pytest.mark.parametrize("reduction", ["mean", "sum"])
+def test_dirichlet_loss_reductions_run(reduction):
+    """Both supported reductions produce a valid scalar.
+
+    Parameters
+    ----------
+    reduction : str
+        The neighbor aggregation reduction to test.
+    """
+    loss_fn = DirichletLoss(reduction=reduction)
+    model_out, batch = _make_inputs()
+    loss = loss_fn.forward(model_out, batch)
+    assert loss.dim() == 0
+    assert torch.isfinite(loss)
+
+
+def test_dirichlet_loss_lambda_scales_linearly():
+    """The output scales linearly with ``lamb``."""
+    model_out, batch = _make_inputs()
+    base = DirichletLoss(lamb=0.1).forward(model_out, batch)
+    scaled = DirichletLoss(lamb=0.5).forward(model_out, batch)
+    assert torch.allclose(scaled, 5.0 * base)
+
+
+def test_dirichlet_loss_zero_when_frames_and_embeddings_align():
+    """Identical embeddings and frames over a cycle give zero loss.
+
+    When every node shares the same embedding and the same frame, each
+    projection is identical, so the neighbor-averaged projection of the current
+    embedding equals the projection of the (equal) initial embedding, and the
+    loss vanishes. The directed cycle guarantees every node has one incoming
+    edge, so no node is left with an all-zero aggregate.
+    """
+    N, r, d = 3, 2, 4
+    shared_emb = torch.ones(N, d)
+    shared_frame = torch.randn(1, r, d).expand(N, r, d).contiguous()
+    model_out = {"x_0": shared_emb, "z_0": shared_emb.clone(), "Q": shared_frame}
+
+    src = torch.arange(N)
+    edge_index = torch.stack([src, torch.roll(src, -1)], dim=0)
+    batch = torch_geometric.data.Data(edge_index=edge_index, num_nodes=N)
+
+    loss = DirichletLoss().forward(model_out, batch)
+    assert torch.allclose(loss, torch.tensor(0.0), atol=1e-6)
+
+
+def test_dirichlet_loss_ignores_self_loops():
+    """Adding a self-loop on every node leaves the loss unchanged.
+
+    The loss strips self-loops from ``edge_index`` before aggregating over
+    neighbors, so an otherwise identical batch augmented with self-loops must
+    produce the same value as the original.
+    """
+    model_out, batch = _make_inputs()
+    looped, _ = torch_geometric.utils.add_self_loops(
+        batch.edge_index, num_nodes=batch.num_nodes
+    )
+    batch_looped = torch_geometric.data.Data(
+        edge_index=looped, num_nodes=batch.num_nodes
+    )
+
+    loss_fn = DirichletLoss()
+    base = loss_fn.forward(model_out, batch)
+    looped_loss = loss_fn.forward(model_out, batch_looped)
+    assert torch.allclose(base, looped_loss, atol=1e-6)
+
+
+def test_dirichlet_loss_detaches_initial_embedding():
+    """Gradients flow to ``x_0`` and ``Q`` but not to the detached ``z_0``."""
+    model_out, batch = _make_inputs(requires_grad=True)
+    loss = DirichletLoss().forward(model_out, batch)
+    loss.backward()
+
+    assert model_out["x_0"].grad is not None
+    assert model_out["Q"].grad is not None
+    # z_0 is used only as a detached target, so no gradient reaches it.
+    assert model_out["z_0"].grad is None
+
+
+def test_dirichlet_loss_one_isolated_node_finite():
+    """A single isolated node among otherwise-connected nodes stays finite."""
+    N = 3
+    edge_index = torch.tensor([[0, 1], [1, 0]], dtype=torch.long)
+    model_out, batch = _make_inputs(N=N, edge_index=edge_index)
+    loss = DirichletLoss().forward(model_out, batch)
+    assert torch.isfinite(loss)
+    assert loss.item() >= 0.0
+
+
+def test_dirichlet_loss_all_nodes_isolated():
+    """A zero-edge graph still produces a finite, non-negative loss."""
+    N = 3
+    edge_index = torch.tensor([[], []], dtype=torch.long)
+    model_out, batch = _make_inputs(N=N, edge_index=edge_index)
+    loss = DirichletLoss().forward(model_out, batch)
+    assert torch.isfinite(loss)
+    assert loss.item() >= 0.0
+
+
+def test_dirichlet_loss_near_zero_projection_grad_finite():
+    """Gradients stay finite when a node's projection is exactly zero."""
+    model_out, batch = _make_inputs(requires_grad=True)
+    with torch.no_grad():
+        model_out["z_0"][0] = 0.0
+    loss = DirichletLoss().forward(model_out, batch)
+    loss.backward()
+    assert torch.isfinite(model_out["Q"].grad).all()
+    assert torch.isfinite(model_out["x_0"].grad).all()
+
+
+def _cycle_edge_index(N, offset=0):
+    """Build a directed cycle over ``N`` nodes, indices shifted by ``offset``.
+
+    Parameters
+    ----------
+    N : int
+        Number of nodes in the cycle.
+    offset : int, optional
+        Amount to shift node indices by, for packing into a larger batch
+        (default: 0).
+
+    Returns
+    -------
+    torch.Tensor
+        Edge index of shape ``[2, N]``.
+    """
+    src = torch.arange(N) + offset
+    dst = torch.roll(src, -1)
+    return torch.stack([src, dst], dim=0)
+
+
+def test_dirichlet_loss_no_cross_graph_leakage():
+    """Perturbing one graph in a batch must not change another graph's gradient.
+
+    Two independent cycle graphs are packed into a single batch, Since the final reduction
+    is a plain mean over all nodes, graph B's gradient should be identical
+    regardless of what graph A's embeddings are.
+    """
+    N_a, N_b, r, d = 3, 3, 2, 4
+    N = N_a + N_b
+    edge_index = torch.cat(
+        [_cycle_edge_index(N_a), _cycle_edge_index(N_b, offset=N_a)], dim=1
+    )
+    batch = torch_geometric.data.Data(edge_index=edge_index, num_nodes=N)
+
+    x_0 = torch.randn(N, d, requires_grad=True)
+    z_0 = torch.randn(N, d)
+    Q = torch.randn(N, r, d, requires_grad=True)
+    model_out = {"x_0": x_0, "z_0": z_0, "Q": Q}
+    loss = DirichletLoss().forward(model_out, batch)
+    loss.backward()
+    grad_b_before = Q.grad[N_a:].clone()
+
+    # Perturb graph A's embeddings only; graph B's data is untouched.
+    x_0b = x_0.detach().clone()
+    x_0b[:N_a] = torch.randn(N_a, d)
+    x_0b.requires_grad_(True)
+    Qb = Q.detach().clone().requires_grad_(True)
+    model_out_2 = {"x_0": x_0b, "z_0": z_0, "Q": Qb}
+    loss2 = DirichletLoss().forward(model_out_2, batch)
+    loss2.backward()
+    grad_b_after = Qb.grad[N_a:]
+
+    assert torch.allclose(grad_b_before, grad_b_after)
+
+
+def test_dirichlet_loss_reductions_actually_differ():
+    """``sum`` and ``mean`` reductions must produce different losses."""
+    N, r, d = 3, 2, 4
+    edge_index = torch.tensor([[0, 1], [2, 2]], dtype=torch.long)
+    model_out, batch = _make_inputs(N=N, r=r, d=d, edge_index=edge_index)
+
+    loss_mean = DirichletLoss(reduction="mean").forward(model_out, batch)
+    loss_sum = DirichletLoss(reduction="sum").forward(model_out, batch)
+
+    assert not torch.allclose(loss_mean, loss_sum)
diff --git a/test/nn/backbones/graph/test_gauge.py b/test/nn/backbones/graph/test_gauge.py
new file mode 100644
index 000000000..d0c9f1006
--- /dev/null
+++ b/test/nn/backbones/graph/test_gauge.py
@@ -0,0 +1,969 @@
+"""Unit tests for the gauge-equivariant graph backbone."""
+
+import math
+
+import pytest
+import torch
+import torch_geometric
+from torch import nn
+
+from topobench.nn.backbones.graph.gauge import (
+    FFBlock,
+    GatedFlatteningLayer,
+    GaugeLayer,
+    GaugeModel,
+    LocalCoordinatesLayer,
+    MultiHeadFF,
+    MultiHeadLinear,
+    NodeUpdateLayer,
+    activation_dict,
+)
+from topobench.nn.wrappers.graph import GaugeWrapper
+
+
+def _is_orthonormal(Q, atol=1e-5):
+    """Check that the frames in ``Q`` have orthonormal rows.
+
+    Parameters
+    ----------
+    Q : torch.Tensor
+        Per-node frames of shape ``[N, r, d]``.
+    atol : float, optional
+        Absolute tolerance for the identity comparison (default: 1e-5).
+
+    Returns
+    -------
+    bool
+        True if ``Q @ Q^T`` equals the ``r x r`` identity for every node.
+    """
+    N, r, _ = Q.shape
+    gram = Q @ Q.transpose(-2, -1)
+    eye = torch.eye(r).expand(N, r, r)
+    return torch.allclose(gram, eye, atol=atol)
+
+
+class TestMultiHeadLinear:
+    """Tests for the per-head linear layer."""
+
+    def test_forward_shape(self):
+        """Output has shape ``[N, r, out_channels]``."""
+        layer = MultiHeadLinear(in_channels=4, out_channels=6, r_dim=3)
+        Z = torch.randn(10, 3, 4)
+        out = layer(Z)
+        assert out.shape == (10, 3, 6)
+
+    def test_parameter_shapes(self):
+        """Stacked weight and bias have the expected per-head shapes."""
+        layer = MultiHeadLinear(in_channels=4, out_channels=6, r_dim=3)
+        assert layer.superW.shape == (3, 6, 4)
+        assert layer.superB.shape == (3, 6)
+
+    def test_no_bias(self):
+        """With ``bias=False`` no bias parameter is registered."""
+        layer = MultiHeadLinear(
+            in_channels=4, out_channels=6, r_dim=3, bias=False
+        )
+        assert layer.superB is None
+        Z = torch.randn(5, 3, 4)
+        assert layer(Z).shape == (5, 3, 6)
+
+    def test_heads_are_independent(self):
+        """Each head applies its own weight matrix and bias.
+
+        The batched ``einsum`` must agree with applying each head's linear map
+        one at a time.
+        """
+        layer = MultiHeadLinear(in_channels=4, out_channels=6, r_dim=3)
+        Z = torch.randn(10, 3, 4)
+        out = layer(Z)
+        for h in range(3):
+            expected = Z[:, h, :] @ layer.superW[h].T + layer.superB[h]
+            assert torch.allclose(out[:, h, :], expected, atol=1e-5)
+
+    def test_permuting_one_head_leaves_others_untouched(self):
+        """Changing a head's input only changes that head's output."""
+        layer = MultiHeadLinear(
+            in_channels=4, out_channels=6, r_dim=3, bias=False
+        )
+        Z = torch.randn(10, 3, 4)
+        out = layer(Z)
+        Z2 = Z.clone()
+        Z2[:, 1, :] = torch.randn(10, 4)
+        out2 = layer(Z2)
+        assert torch.allclose(out[:, 0, :], out2[:, 0, :])
+        assert torch.allclose(out[:, 2, :], out2[:, 2, :])
+        assert not torch.allclose(out[:, 1, :], out2[:, 1, :])
+
+    def test_reset_parameters_bound(self):
+        """Weights are initialized within ``1 / sqrt(in_channels)``."""
+        in_channels = 9
+        layer = MultiHeadLinear(
+            in_channels=in_channels, out_channels=6, r_dim=3
+        )
+        bound = 1 / math.sqrt(in_channels)
+        assert layer.superW.abs().max().item() <= bound + 1e-6
+        assert layer.superB.abs().max().item() <= bound + 1e-6
+
+    def test_reset_parameters_changes_weights(self):
+        """Calling ``reset_parameters`` re-samples the weights."""
+        layer = MultiHeadLinear(in_channels=4, out_channels=6, r_dim=3)
+        before = layer.superW.clone()
+        layer.reset_parameters()
+        assert not torch.allclose(before, layer.superW)
+
+
+class TestMultiHeadFF:
+    """Tests for the per-head feed-forward network."""
+
+    def test_single_layer_shape(self):
+        """With ``hidden_dims=None`` the network is a single per-head map."""
+        net = MultiHeadFF(in_channels=4, out_channels=2, r=3)
+        # exactly one MultiHeadLinear, no activation / dropout
+        assert len(net.model) == 1
+        assert isinstance(net.model[0], MultiHeadLinear)
+        Z = torch.randn(7, 3, 4)
+        assert net(Z).shape == (7, 3, 2)
+
+    def test_multi_layer_shape(self):
+        """Hidden dims add intermediate per-head layers."""
+        net = MultiHeadFF(
+            in_channels=4, out_channels=2, r=3, hidden_dims=[8, 8]
+        )
+        Z = torch.randn(7, 3, 4)
+        assert net(Z).shape == (7, 3, 2)
+
+    def test_no_activation_after_output(self):
+        """Activation and dropout appear only between layers."""
+        net = MultiHeadFF(in_channels=4, out_channels=2, r=3, hidden_dims=[8])
+        # Sequence: Linear, Act, Dropout, Linear -> last module is Linear
+        assert isinstance(net.model[-1], MultiHeadLinear)
+        linears = [m for m in net.model if isinstance(m, MultiHeadLinear)]
+        assert len(linears) == 2
+
+    @pytest.mark.parametrize("act", list(activation_dict.keys()))
+    def test_activation_choice(self, act):
+        """Every activation in ``activation_dict`` is wired up correctly.
+
+        Parameters
+        ----------
+        act : str
+            Name of the activation function to test.
+        """
+        net = MultiHeadFF(
+            in_channels=4, out_channels=2, r=3, hidden_dims=[8], act=act
+        )
+        acts = [m for m in net.model if isinstance(m, activation_dict[act])]
+        assert len(acts) == 1
+        Z = torch.randn(7, 3, 4)
+        assert net(Z).shape == (7, 3, 2)
+
+    def test_final_activation(self):
+        """``final_activation`` appends the activation after the output layer."""
+        net = MultiHeadFF(
+            in_channels=4,
+            out_channels=2,
+            r=3,
+            hidden_dims=[8],
+            act="leaky_relu",
+            final_activation=True,
+        )
+        # Output layer is followed by the activation, not left bare.
+        assert isinstance(net.model[-1], nn.LeakyReLU)
+        # Two activations total: one between layers, one after the output.
+        acts = [m for m in net.model if isinstance(m, nn.LeakyReLU)]
+        assert len(acts) == 2
+        Z = torch.randn(7, 3, 4)
+        assert net(Z).shape == (7, 3, 2)
+
+
+class TestFFBlock:
+    """Tests for the feed-forward block."""
+
+    def test_forward_shape(self):
+        """Output has the requested number of channels."""
+        block = FFBlock(in_channels=4, out_channels=8, hidden_dim=16)
+        x = torch.randn(5, 4)
+        assert block(x).shape == (5, 8)
+
+    def test_norm_on_input_channels(self):
+        """The pre-norm LayerNorm normalizes over ``in_channels``."""
+        block = FFBlock(in_channels=4, out_channels=8, hidden_dim=16)
+        assert block.norm.normalized_shape == (4,)
+
+    def test_extra_leading_dims(self):
+        """The block broadcasts over arbitrary leading dimensions."""
+        block = FFBlock(in_channels=4, out_channels=8, hidden_dim=16)
+        x = torch.randn(5, 3, 4)
+        assert block(x).shape == (5, 3, 8)
+
+    def test_rejects_negative_hidden_layers(self):
+        """A negative number of hidden layers is rejected."""
+        with pytest.raises(AssertionError):
+            FFBlock(
+                in_channels=4, out_channels=8, hidden_dim=16, n_hidden_layers=-1
+            )
+
+    def test_zero_hidden_layers_is_single_linear(self):
+        """With no hidden layers the block is a single linear map."""
+        block = FFBlock(
+            in_channels=4, out_channels=8, hidden_dim=16, n_hidden_layers=0
+        )
+        linears = [m for m in block.model if isinstance(m, nn.Linear)]
+        assert len(linears) == 1
+        assert linears[0].in_features == 4
+        assert linears[0].out_features == 8
+        assert block(torch.randn(5, 4)).shape == (5, 8)
+
+    @pytest.mark.parametrize("n_hidden_layers", [1, 2, 3])
+    def test_num_hidden_layers(self, n_hidden_layers):
+        """Multiple hidden layers still produce the right output shape.
+
+        Parameters
+        ----------
+        n_hidden_layers : int
+            Number of hidden layers to configure.
+        """
+        block = FFBlock(
+            in_channels=4,
+            out_channels=8,
+            hidden_dim=16,
+            n_hidden_layers=n_hidden_layers,
+        )
+        x = torch.randn(5, 4)
+        assert block(x).shape == (5, 8)
+
+    def test_default_activation_is_gelu(self):
+        """The default activation stays GELU for backward compatibility."""
+        block = FFBlock(in_channels=4, out_channels=8, hidden_dim=16)
+        acts = [m for m in block.model if isinstance(m, nn.GELU)]
+        assert len(acts) == 1
+
+    @pytest.mark.parametrize("act", list(activation_dict.keys()))
+    def test_activation_choice(self, act):
+        """Every activation in ``activation_dict`` is wired up correctly.
+
+        Parameters
+        ----------
+        act : str
+            Name of the activation function to test.
+        """
+        block = FFBlock(
+            in_channels=4,
+            out_channels=8,
+            hidden_dim=16,
+            n_hidden_layers=2,
+            act=act,
+        )
+        acts = [m for m in block.model if isinstance(m, activation_dict[act])]
+        # one activation per hidden layer
+        assert len(acts) == 2
+        x = torch.randn(5, 4)
+        assert block(x).shape == (5, 8)
+
+
+class TestLocalCoordinatesLayer:
+    """Tests for the local coordinate (frame) layer."""
+
+    def test_output_shape(self, simple_graph_0):
+        """The layer returns one ``[r, d]`` frame per node.
+
+        Parameters
+        ----------
+        simple_graph_0 : torch_geometric.data.Data
+            Test graph fixture.
+        """
+        d, r = 8, 3
+        layer = LocalCoordinatesLayer(r_subspaces=r, d_embedd=d)
+        Z = torch.randn(simple_graph_0.num_nodes, d)
+        Q = layer(Z, simple_graph_0.edge_index)
+        assert Q.shape == (simple_graph_0.num_nodes, r, d)
+
+    def test_frames_orthonormal(self, simple_graph_0):
+        """The QR step yields orthonormal per-node frames.
+
+        Parameters
+        ----------
+        simple_graph_0 : torch_geometric.data.Data
+            Test graph fixture.
+        """
+        d, r = 8, 3
+        layer = LocalCoordinatesLayer(r_subspaces=r, d_embedd=d)
+        Z = torch.randn(simple_graph_0.num_nodes, d)
+        Q = layer(Z, simple_graph_0.edge_index)
+        assert _is_orthonormal(Q)
+
+    def test_no_nan(self, simple_graph_0):
+        """Output frames are finite.
+
+        Parameters
+        ----------
+        simple_graph_0 : torch_geometric.data.Data
+            Test graph fixture.
+        """
+        d, r = 8, 4
+        layer = LocalCoordinatesLayer(r_subspaces=r, d_embedd=d)
+        Z = torch.randn(simple_graph_0.num_nodes, d)
+        Q = layer(Z, simple_graph_0.edge_index)
+        assert not torch.isnan(Q).any()
+        assert not torch.isinf(Q).any()
+
+    def test_all_nodes_isolated_no_nan(self):
+        """
+        A graph with zero edges still produces finite, orthonormal frames.
+
+        """
+        d, r = 8, 3
+        edge_index = torch.tensor([[], []], dtype=torch.long)
+        Z = torch.randn(3, d)
+        layer = LocalCoordinatesLayer(r_subspaces=r, d_embedd=d)
+        Q = layer(Z, edge_index)
+        assert Q.shape == (3, r, d)
+        assert _is_orthonormal(Q)
+        assert not torch.isnan(Q).any()
+
+    def test_isolated_nodes_independent(self):
+        """
+        With no edges, changing one node's features must not affect others.
+
+        """
+        d, r = 8, 3
+        edge_index = torch.tensor([[],[]], dtype =torch.long)
+        Z = torch.randn(3, d)
+        Z1 = Z.clone()
+        Z1[1] = torch.randn(d)
+        layer = LocalCoordinatesLayer(r_subspaces=r, d_embedd=d)
+        layer.eval()
+        Q = layer(Z, edge_index)
+        Q1 = layer(Z1, edge_index)
+        assert torch.allclose(Q[0], Q1[0])
+        assert torch.allclose(Q[2], Q1[2])
+        assert not torch.allclose(Q[1], Q1[1])
+
+    def test_r_greater_than_d_raises(self):
+        """``r > d_embedd`` is rejected instead of silently truncating to d.
+
+        A d-dimensional space admits at most d orthonormal frame vectors, so
+        the layer must refuse to be constructed rather than dropping subspaces.
+        """
+        with pytest.raises(ValueError):
+            LocalCoordinatesLayer(r_subspaces=8, d_embedd=3)
+
+    def test_respects_edge_direction(self):
+        """A directed edge only influences its destination, never its source."""
+        d, r = 8, 3
+        edge_index = torch.tensor([[0], [1]], dtype=torch.long)
+        Z = torch.randn(2, d)
+        Z1 = Z.clone()
+        Z1[1] = torch.randn(d)
+        layer = LocalCoordinatesLayer(r_subspaces=r, d_embedd=d)
+        layer.eval()
+        Q = layer(Z, edge_index)
+        Q1 = layer(Z1, edge_index)
+        assert torch.allclose(Q[0], Q1[0])
+        assert not torch.allclose(Q[1], Q1[1])
+
+    def test_star_topology_no_nan(self):
+        """A hub node with far more neighbors than the rest of the graph must not produce NaNs."""
+        d, r = 8, 3
+        n_leaves = 20
+        N = n_leaves + 1
+        leaves = torch.arange(1, N)
+        hub = torch.zeros(n_leaves, dtype=torch.long)
+        src = torch.cat([hub, leaves])
+        dst = torch.cat([leaves, hub])
+        edge_index = torch.stack([src, dst])
+        layer = LocalCoordinatesLayer(r_subspaces=r, d_embedd=d)
+        Z = torch.randn(N, d)
+        Q = layer(Z, edge_index)
+        assert Q.shape == (N, r, d)
+        assert _is_orthonormal(Q)
+        assert not torch.isnan(Q).any() and not torch.isinf(Q).any()
+
+
+class TestGatedFlatteningLayer:
+    """Tests for the gated flattening (frame smoothing) layer."""
+
+    def test_output_shape_and_orthonormality(self, simple_graph_0):
+        """Smoothing preserves the shape and orthonormality of the frames.
+
+        Parameters
+        ----------
+        simple_graph_0 : torch_geometric.data.Data
+            Test graph fixture.
+        """
+        d, r = 8, 3
+        N = simple_graph_0.num_nodes
+        # Build orthonormal input frames via the local-coords layer.
+        coords = LocalCoordinatesLayer(r_subspaces=r, d_embedd=d)
+        Q = coords(torch.randn(N, d), simple_graph_0.edge_index)
+
+        gate = GatedFlatteningLayer(r=r)
+        Qnew = gate(Q, simple_graph_0.edge_index)
+        assert Qnew.shape == (N, r, d)
+        assert _is_orthonormal(Qnew)
+
+    def test_no_learnable_parameters(self):
+        """The gated flattening layer is parameter-free."""
+        gate = GatedFlatteningLayer(r=3)
+        assert list(gate.parameters()) == []
+
+    @pytest.mark.xfail(
+        reason="QR's backward is undefined for rank-deficient input; "
+        "gamma=1.0 zeroes an isolated node's blend, causing NaN/Inf grads."
+    )
+    def test_isolated_node_gamma_one_backward(self):
+        """Gradients must stay finite even when gamma=1.0 zeroes an isolated node's blend."""
+        d, r = 8, 3
+        N = 3
+        edge_index = torch.tensor([[], []], dtype=torch.long)
+        Z = torch.randn(N, d, requires_grad=True)
+        Q = LocalCoordinatesLayer(r_subspaces=r, d_embedd=d)(Z, edge_index)
+        gate = GatedFlatteningLayer(r=r, gamma=1.0)
+        Qnew = gate(Q, edge_index)
+        Qnew.sum().backward()
+        assert Z.grad is not None
+        assert not torch.isnan(Z.grad).any()
+        assert not torch.isinf(Z.grad).any()
+
+
+class TestNodeUpdateLayer:
+    """Tests for the node feature update layer."""
+
+    def test_output_shape(self, simple_graph_0):
+        """The update returns ``out_channels`` features per node.
+
+        Parameters
+        ----------
+        simple_graph_0 : torch_geometric.data.Data
+            Test graph fixture.
+        """
+        d, r = 8, 3
+        N = simple_graph_0.num_nodes
+        layer = NodeUpdateLayer(in_channels=d, out_channels=d)
+        Z = torch.randn(N, d)
+        Q = LocalCoordinatesLayer(r_subspaces=r, d_embedd=d)(
+            Z, simple_graph_0.edge_index
+        )
+        out = layer(Z, Q, simple_graph_0.edge_index)
+        assert out.shape == (N, d)
+
+    def test_residual_disabled(self, simple_graph_0):
+        """``phi_hidden_layers=None`` disables the residual MLP.
+
+        Parameters
+        ----------
+        simple_graph_0 : torch_geometric.data.Data
+            Test graph fixture.
+        """
+        d, r = 8, 3
+        N = simple_graph_0.num_nodes
+        layer = NodeUpdateLayer(
+            in_channels=d, out_channels=d, phi_hidden_layers=None
+        )
+        assert layer.phi is None
+        Z = torch.randn(N, d)
+        Q = LocalCoordinatesLayer(r_subspaces=r, d_embedd=d)(
+            Z, simple_graph_0.edge_index
+        )
+        assert layer(Z, Q, simple_graph_0.edge_index).shape == (N, d)
+
+    def test_residual_enabled(self, simple_graph_0):
+        """The residual MLP is built when enabled.
+
+        Parameters
+        ----------
+        simple_graph_0 : torch_geometric.data.Data
+            Test graph fixture.
+        """
+        layer = NodeUpdateLayer(
+            in_channels=8, out_channels=8, phi_hidden_layers=1
+        )
+        assert isinstance(layer.phi, FFBlock)
+
+    def test_supports_different_in_out_channels(self, simple_graph_0):
+        """The update maps to ``out_channels`` even when it differs from ``in_channels``.
+
+        Parameters
+        ----------
+        simple_graph_0 : torch_geometric.data.Data
+            Test graph fixture.
+        """
+        d, r = 8, 3
+        out_channels = 12
+        N = simple_graph_0.num_nodes
+        Z = torch.randn(N, d)
+        Q = LocalCoordinatesLayer(r_subspaces=r, d_embedd=d)(
+            Z, simple_graph_0.edge_index
+        )
+        layer = NodeUpdateLayer(
+            in_channels=d, out_channels=out_channels, phi_hidden_layers=1
+        )
+        assert layer(Z, Q, simple_graph_0.edge_index).shape == (N, out_channels)
+
+
+class TestGaugeLayer:
+    """Tests for a single gauge message-passing layer."""
+
+    def test_output_shapes(self, simple_graph_0):
+        """The layer returns updated features and orthonormal frames.
+
+        Parameters
+        ----------
+        simple_graph_0 : torch_geometric.data.Data
+            Test graph fixture.
+        """
+        d, r = 8, 3
+        N = simple_graph_0.num_nodes
+        layer = GaugeLayer(d_embedd=d, r=r, n_gated=2)
+        x = torch.randn(N, d)
+        Znew, Q = layer(x, simple_graph_0.edge_index)
+        assert Znew.shape == (N, d)
+        assert Q.shape == (N, r, d)
+        assert _is_orthonormal(Q)
+
+    @pytest.mark.parametrize("n_gated", [0, 1, 3])
+    def test_num_gated_layers(self, simple_graph_0, n_gated):
+        """The number of gated flattening sublayers is configurable.
+
+        Parameters
+        ----------
+        simple_graph_0 : torch_geometric.data.Data
+            Test graph fixture.
+        n_gated : int
+            Number of gated flattening layers to stack.
+        """
+        d, r = 8, 3
+        layer = GaugeLayer(d_embedd=d, r=r, n_gated=n_gated)
+        assert len(layer.gated_flattening_layers) == n_gated
+        x = torch.randn(simple_graph_0.num_nodes, d)
+        Znew, Q = layer(x, simple_graph_0.edge_index)
+        assert _is_orthonormal(Q)
+
+    def test_all_nodes_isolated_no_nan(self):
+
+        """A zero-edge graph stays finite through the full gauge layer."""
+
+        d, r = 8, 3
+        N = 3
+        edge_index = torch.tensor([[], []], dtype=torch.long)
+        layer = GaugeLayer(d_embedd=d, r=r, n_gated=2)
+        x = torch.randn(N, d)
+        Znew, Q = layer(x, edge_index)
+        assert Znew.shape == (N, d)
+        assert Q.shape == (N, r, d)
+        assert _is_orthonormal(Q)
+        assert not torch.isnan(Znew).any() and not torch.isinf(Znew).any()
+        assert not torch.isnan(Q).any() and not torch.isinf(Q).any()
+
+
+
+class TestGaugeModel:
+    """Tests for the full gauge model."""
+
+    def test_forward_shapes(self, simple_graph_0):
+        """The model maps input features to embeddings and final frames.
+
+        Parameters
+        ----------
+        simple_graph_0 : torch_geometric.data.Data
+            Test graph fixture.
+        """
+        in_channels, d, r = 5, 8, 3
+        N = simple_graph_0.num_nodes
+        model = GaugeModel(
+            n_layers=2, in_channels=in_channels, r=r, d_embedd=d
+        )
+        x = torch.randn(N, in_channels)
+        z, Q = model(x, simple_graph_0.edge_index)
+        assert z.shape == (N, d)
+        assert Q.shape == (N, r, d)
+        assert _is_orthonormal(Q)
+
+    def test_self_loops_do_not_affect_output(self, simple_graph_0):
+        """Adding self-loops leaves the model output unchanged.
+
+        The model strips self-loops from ``edge_index`` before aggregating, so
+        feeding a graph augmented with a self-loop on every node must yield the
+        same embeddings and frames as the original graph.
+
+        Parameters
+        ----------
+        simple_graph_0 : torch_geometric.data.Data
+            Test graph fixture.
+        """
+        in_channels, d, r = 5, 8, 3
+        N = simple_graph_0.num_nodes
+        model = GaugeModel(
+            n_layers=2, in_channels=in_channels, r=r, d_embedd=d
+        )
+        model.eval()  # disable dropout so the comparison is deterministic
+        x = torch.randn(N, in_channels)
+
+        edge_index = simple_graph_0.edge_index
+        looped, _ = torch_geometric.utils.add_self_loops(
+            edge_index, num_nodes=N
+        )
+
+        z, Q = model(x, edge_index)
+        z_looped, Q_looped = model(x, looped)
+
+        assert torch.allclose(z, z_looped, atol=1e-6)
+        assert torch.allclose(Q, Q_looped, atol=1e-6)
+
+    def test_return_initial_true_returns_initial_projection(
+        self, simple_graph_0
+    ):
+        """With ``return_initial=True`` the initial projection is returned.
+
+        The default forward pass yields the triple ``(z, Q, z0)`` where ``z0``
+        is exactly the input projection of ``x`` (before any gauge layer).
+
+        Parameters
+        ----------
+        simple_graph_0 : torch_geometric.data.Data
+            Test graph fixture.
+        """
+        in_channels, d, r = 5, 8, 3
+        N = simple_graph_0.num_nodes
+        model = GaugeModel(
+            n_layers=2, in_channels=in_channels, r=r, d_embedd=d
+        )
+        model.eval()
+        x = torch.randn(N, in_channels)
+
+        out = model(x, simple_graph_0.edge_index, return_initial=True)
+        assert len(out) == 3
+        z, Q, z0 = out
+        assert z.shape == (N, d)
+        assert Q.shape == (N, r, d)
+        assert z0.shape == (N, d)
+        # z0 is precisely the input projection, unaffected by the gauge layers.
+        assert torch.allclose(z0, model.input_projector(x))
+
+    def test_return_initial_default_returns_pair(self, simple_graph_0):
+        """Omitting ``return_initial`` returns the ``(z, Q)`` pair by default.
+
+        Parameters
+        ----------
+        simple_graph_0 : torch_geometric.data.Data
+            Test graph fixture.
+        """
+        model = GaugeModel(n_layers=2, in_channels=5, r=3, d_embedd=8)
+        x = torch.randn(simple_graph_0.num_nodes, 5)
+        assert len(model(x, simple_graph_0.edge_index)) == 2
+
+    def test_return_initial_false_returns_pair(self, simple_graph_0):
+        """With ``return_initial=False`` only ``(z, Q)`` is returned.
+
+        Parameters
+        ----------
+        simple_graph_0 : torch_geometric.data.Data
+            Test graph fixture.
+        """
+        model = GaugeModel(n_layers=2, in_channels=5, r=3, d_embedd=8)
+        x = torch.randn(simple_graph_0.num_nodes, 5)
+        out = model(x, simple_graph_0.edge_index)
+        assert len(out) == 2
+
+    @pytest.mark.parametrize("n_layers", [1, 2, 4])
+    def test_num_layers(self, simple_graph_0, n_layers):
+        """The model stacks the requested number of gauge layers.
+
+        Parameters
+        ----------
+        simple_graph_0 : torch_geometric.data.Data
+            Test graph fixture.
+        n_layers : int
+            Number of gauge layers to stack.
+        """
+        in_channels, d, r = 5, 8, 3
+        model = GaugeModel(
+            n_layers=n_layers, in_channels=in_channels, r=r, d_embedd=d
+        )
+        assert len(model.layers) == n_layers
+        x = torch.randn(simple_graph_0.num_nodes, in_channels)
+        z, _ = model(x, simple_graph_0.edge_index)
+        assert z.shape == (simple_graph_0.num_nodes, d)
+
+    def test_no_nan(self, simple_graph_0):
+        """The final embeddings are finite.
+
+        Parameters
+        ----------
+        simple_graph_0 : torch_geometric.data.Data
+            Test graph fixture.
+        """
+        model = GaugeModel(n_layers=2, in_channels=5, r=3, d_embedd=8)
+        x = torch.randn(simple_graph_0.num_nodes, 5)
+        z, Q = model(x, simple_graph_0.edge_index)
+        assert not torch.isnan(z).any() and not torch.isinf(z).any()
+        assert not torch.isnan(Q).any() and not torch.isinf(Q).any()
+
+    def test_backward_pass(self, simple_graph_0):
+        """Gradients flow back to the input and the parameters.
+
+        Parameters
+        ----------
+        simple_graph_0 : torch_geometric.data.Data
+            Test graph fixture.
+        """
+        model = GaugeModel(n_layers=2, in_channels=5, r=3, d_embedd=8)
+        x = torch.randn(simple_graph_0.num_nodes, 5, requires_grad=True)
+        z, _ = model(x, simple_graph_0.edge_index)
+        z.sum().backward()
+        assert x.grad is not None
+        has_grad = any(
+            p.grad is not None and p.grad.abs().sum() > 0
+            for p in model.parameters()
+            if p.requires_grad
+        )
+        assert has_grad
+
+    def test_deterministic_in_eval(self, simple_graph_0):
+        """Two eval-mode passes on the same input agree.
+
+        Parameters
+        ----------
+        simple_graph_0 : torch_geometric.data.Data
+            Test graph fixture.
+        """
+        model = GaugeModel(n_layers=2, in_channels=5, r=3, d_embedd=8)
+        model.eval()
+        x = torch.randn(simple_graph_0.num_nodes, 5)
+        z1, _ = model(x, simple_graph_0.edge_index)
+        z2, _ = model(x, simple_graph_0.edge_index)
+        assert torch.allclose(z1, z2)
+
+    def test_batched_graphs(self, simple_graph_0, simple_graph_1):
+        """The model handles a batch of disconnected graphs.
+
+        Parameters
+        ----------
+        simple_graph_0 : torch_geometric.data.Data
+            First test graph fixture.
+        simple_graph_1 : torch_geometric.data.Data
+            Second test graph fixture.
+        """
+        batch = torch_geometric.data.Batch.from_data_list(
+            [simple_graph_0, simple_graph_1]
+        )
+        n_total = simple_graph_0.num_nodes + simple_graph_1.num_nodes
+        model = GaugeModel(n_layers=2, in_channels=5, r=3, d_embedd=8)
+        x = torch.randn(n_total, 5)
+        z, _ = model(x, batch.edge_index)
+        assert z.shape == (n_total, 8)
+
+    @pytest.mark.parametrize("act", list(activation_dict.keys()))
+    def test_activation_propagates_to_ffblocks(self, act):
+        """The ``act`` argument reaches every feed-forward block.
+
+        The activation must flow from the model down to the ``fflayer`` of each
+        local-coordinates layer and the ``phi`` residual of each node update.
+
+        Parameters
+        ----------
+        act : str
+            Name of the activation function to test.
+        """
+        model = GaugeModel(n_layers=2, in_channels=5, r=3, d_embedd=8, act=act)
+        expected = activation_dict[act]
+        for layer in model.layers:
+            fflayer = layer.local_coords_layer.fflayer
+            assert any(isinstance(m, expected) for m in fflayer.model)
+            phi = layer.node_update_layer.phi
+            assert any(isinstance(m, expected) for m in phi.model)
+
+    def test_f_sim_act_defaults_to_leaky_relu(self):
+        """The similarity scorer defaults to LeakyReLU, independent of ``act``."""
+        model = GaugeModel(
+            n_layers=2, in_channels=5, r=3, d_embedd=8, act="gelu"
+        )
+        for layer in model.layers:
+            f_sim = layer.local_coords_layer.f_sim
+            assert any(isinstance(m, nn.LeakyReLU) for m in f_sim.model)
+
+    @pytest.mark.parametrize("f_sim_act", list(activation_dict.keys()))
+    def test_f_sim_act_applied_to_final_score(self, f_sim_act):
+        """``f_sim`` ends on its activation, matching the reference score_lin.
+
+        The reference ``score_lin`` (Linear -> activation) activates the score
+        before the softmax; here ``f_sim`` is built with ``final_activation``,
+        so the last module of its sequential is the activation.
+
+        Parameters
+        ----------
+        f_sim_act : str
+            Name of the similarity-scorer activation to test.
+        """
+        model = GaugeModel(
+            n_layers=2,
+            in_channels=5,
+            r=3,
+            d_embedd=8,
+            f_sim_act=f_sim_act,
+        )
+        expected = activation_dict[f_sim_act]
+        for layer in model.layers:
+            f_sim = layer.local_coords_layer.f_sim
+            assert f_sim.final_activation is True
+            assert isinstance(f_sim.model[-1], expected)
+
+    @pytest.mark.parametrize("f_sim_act", list(activation_dict.keys()))
+    def test_f_sim_act_propagates(self, f_sim_act):
+        """The ``f_sim_act`` argument reaches every similarity network.
+
+        The knob is independent of ``act``: it must only affect ``f_sim``, not
+        the ``fflayer`` feed-forward blocks.
+
+        Parameters
+        ----------
+        f_sim_act : str
+            Name of the similarity-scorer activation to test.
+        """
+        model = GaugeModel(
+            n_layers=2,
+            in_channels=5,
+            r=3,
+            d_embedd=8,
+            act="gelu",
+            f_sim_act=f_sim_act,
+        )
+        expected = activation_dict[f_sim_act]
+        for layer in model.layers:
+            f_sim = layer.local_coords_layer.f_sim
+            assert any(isinstance(m, expected) for m in f_sim.model)
+            # ``act`` still governs the feed-forward block independently.
+            fflayer = layer.local_coords_layer.fflayer
+            assert any(isinstance(m, nn.GELU) for m in fflayer.model)
+
+    def test_dropout_propagates_to_ffblocks(self):
+        """The ``dropout`` argument reaches every feed-forward block.
+
+        The probability must flow from the model down to the ``fflayer`` of
+        each local-coordinates layer and the ``phi`` residual of each node
+        update.
+        """
+        model = GaugeModel(
+            n_layers=2, in_channels=5, r=3, d_embedd=8, dropout=0.42
+        )
+        for layer in model.layers:
+            fflayer = layer.local_coords_layer.fflayer
+            phi = layer.node_update_layer.phi
+            for block in (fflayer, phi):
+                dropouts = [
+                    m for m in block.model if isinstance(m, nn.Dropout)
+                ]
+                assert dropouts
+                assert all(m.p == 0.42 for m in dropouts)
+
+    def test_f_sim_dropout_propagates(self):
+        """The ``f_sim_dropout`` knob reaches ``f_sim`` and is independent.
+
+        It must only configure the similarity network, leaving the
+        feed-forward block dropout governed by ``dropout``. Because ``f_sim``
+        has no hidden layers, ``MultiHeadFF`` (which applies dropout only
+        between layers) instantiates no active dropout module, so the check is
+        on the propagated probability rather than on the module list.
+        """
+        model = GaugeModel(
+            n_layers=2,
+            in_channels=5,
+            r=3,
+            d_embedd=8,
+            dropout=0.1,
+            f_sim_dropout=0.5,
+        )
+        for layer in model.layers:
+            local = layer.local_coords_layer
+            assert local.f_sim_dropout == 0.5
+            assert local.f_sim.dropout == 0.5
+            # The feed-forward block keeps its own dropout probability.
+            ff_drops = [
+                m for m in local.fflayer.model if isinstance(m, nn.Dropout)
+            ]
+            assert all(m.p == 0.1 for m in ff_drops)
+
+    def test_all_nodes_isolated_no_nan(self):
+        """A zero-edge graph stays finite through the full stacked model."""
+        in_channels, d, r, N = 5, 8, 3, 3
+        edge_index = torch.tensor([[], []], dtype=torch.long)
+        model = GaugeModel(
+            n_layers=2, in_channels=in_channels, r=r, d_embedd=d
+        )
+        x = torch.randn(N, in_channels)
+        z, Q = model(x, edge_index)
+        assert z.shape == (N, d)
+        assert Q.shape == (N, r, d)
+        assert _is_orthonormal(Q)
+        assert not torch.isnan(z).any() and not torch.isinf(z).any()
+        assert not torch.isnan(Q).any() and not torch.isinf(Q).any()
+
+
+class TestGaugeWrapper:
+    """Tests for the topobench wrapper around the gauge model."""
+
+    def test_forward(self, simple_graph_0):
+        """The wrapper forwards node embeddings and the initial state.
+
+        In addition to ``x_0`` it exposes the initial projection ``z_0`` and
+        the final frames ``Q`` (consumed by the custom loss).
+
+        Parameters
+        ----------
+        simple_graph_0 : torch_geometric.data.Data
+            Test graph fixture.
+        """
+        # The wrapper adds a residual (batch.x_0 + model output), so the input
+        # width must match the embedding width ``d``.
+        in_channels = d = 8
+        r = 3
+        N = simple_graph_0.num_nodes
+        model = GaugeModel(
+            n_layers=2, in_channels=in_channels, r=r, d_embedd=d
+        )
+        wrapper = GaugeWrapper(model, out_channels=d, num_cell_dimensions=1)
+        batch = torch_geometric.data.Data(
+            x_0=torch.randn(N, in_channels),
+            edge_index=simple_graph_0.edge_index,
+            y=simple_graph_0.y,
+            batch_0=torch.zeros(N, dtype=torch.long),
+        )
+        model_out = wrapper(batch)
+        assert model_out["x_0"].shape == (N, d)
+        assert model_out["z_0"].shape == (N, d)
+        assert model_out["Q"].shape == (N, r, d)
+        assert "labels" in model_out
+        assert "batch_0" in model_out
+
+    def test_forward_residual_off(self, simple_graph_0):
+        """With the residual disabled the backbone width may differ from the input.
+
+        This mirrors the shipped config (``residual_connections: false``,
+        ``out_channels = d_embedd`` != encoder width), so the embedding width is
+        free to differ from ``in_channels``.
+
+        Parameters
+        ----------
+        simple_graph_0 : torch_geometric.data.Data
+            Test graph fixture.
+        """
+        in_channels, d, r = 5, 8, 3
+        N = simple_graph_0.num_nodes
+        model = GaugeModel(
+            n_layers=2, in_channels=in_channels, r=r, d_embedd=d
+        )
+        wrapper = GaugeWrapper(
+            model,
+            out_channels=d,
+            num_cell_dimensions=1,
+            residual_connections=False,
+        )
+        assert wrapper.residual_connections is False
+        batch = torch_geometric.data.Data(
+            x_0=torch.randn(N, in_channels),
+            edge_index=simple_graph_0.edge_index,
+            y=simple_graph_0.y,
+            batch_0=torch.zeros(N, dtype=torch.long),
+        )
+        model_out = wrapper(batch)
+        assert model_out["x_0"].shape == (N, d)
diff --git a/test/pipeline/test_pipeline.py b/test/pipeline/test_pipeline.py
index a61165ae9..d7179002e 100644
--- a/test/pipeline/test_pipeline.py
+++ b/test/pipeline/test_pipeline.py
@@ -7,7 +7,7 @@
 
 
 DATASET = "graph/MUTAG"  # ADD YOUR DATASET HERE
-MODELS = ["graph/gcn", "cell/topotune", "simplicial/topotune"]  # ADD ONE OR SEVERAL MODELS
+MODELS = ["graph/gauge", "graph/gcn", "cell/topotune", "simplicial/topotune"]  # ADD ONE OR SEVERAL MODELS
 
 
 class TestPipeline:
diff --git a/topobench/loss/model/DirichletLoss.py b/topobench/loss/model/DirichletLoss.py
new file mode 100644
index 000000000..0bc2874a5
--- /dev/null
+++ b/topobench/loss/model/DirichletLoss.py
@@ -0,0 +1,86 @@
+"""Dirichlet-energy regularization loss for the Gauge model."""
+
+import torch
+from torch import Tensor
+from torch.nn import functional as F
+from torch_geometric.data import Data
+from torch_geometric.utils import remove_self_loops
+from torch_scatter import scatter
+
+from topobench.loss.base import AbstractLoss
+
+
+class DirichletLoss(AbstractLoss):
+    r"""Dirichlet-energy regularization loss for the Gauge model (cf. paper equation (13)).
+
+    Measures how smoothly the node embeddings vary across edges once projected
+    onto each node's local frame ``Q``. Each node embedding is projected onto
+    its ``r`` frame vectors and L2-normalized; the projection of the current
+    embedding is then averaged over neighbors and compared to the projection of
+    the (detached) initial embedding. The resulting term is scaled by ``lamb``
+    and added to the task loss as a regularizer.
+
+    Parameters
+    ----------
+    lamb : float, optional
+        Weight (lambda) of the regularizer (default: 0.1).
+    reduction : str, optional
+        Neighbor aggregation reduction, either "mean" or "sum" (default: "mean").
+    """
+
+    def __init__(self, lamb: float = 0.1, reduction: str = "mean"):
+        super().__init__()
+
+        if reduction not in ["mean", "sum"]:
+            raise NotImplementedError(
+                f"reduction '{reduction}' not implemented. Valid choices are 'mean', 'sum'."
+            )
+
+        self.lamb = lamb
+        self.reduce = reduction
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}(lamb={self.lamb}, reduction={self.reduce})"
+
+    def forward(self, model_out: dict, batch: Data) -> Tensor:
+        r"""Compute the Dirichlet-energy regularization loss according to the paper's eq. 13.
+
+        Parameters
+        ----------
+        model_out : dict
+            Dictionary containing the model output. Uses ``x_0`` (the final node
+            embeddings ``[N, d]``), ``z_0`` (the initial node embeddings
+            ``[N, d]``) and ``Q`` (the per-node frames ``[N, r, d]``).
+        batch : torch_geometric.data.Data
+            Batch object containing the batched domain data. Uses
+            ``edge_index`` for the neighbor aggregation.
+
+        Returns
+        -------
+        Tensor
+            Scalar regularization loss scaled by ``lamb``.
+        """
+
+        zL = model_out["x_0"]  # [N, d]
+        Q = model_out["Q"]  # Q has shape [N, r, d]
+        z0 = model_out["z_0"]  # [N, d]
+
+        N = z0.size(0)
+        # Removed here (and, consistently, in the Gauge backbone) so the loss
+        # aggregates over the same self-loop-free neighborhoods as the model.
+        edge_index, _ = remove_self_loops(batch.edge_index)
+        src, dst = edge_index[0], edge_index[1]
+
+        # this is essentially the zhat=StopGrad(z0)
+        zhat = z0.detach()
+
+        src_term = F.normalize(torch.einsum("ijk,ik->ij", Q, zhat), dim=-1)
+        agg_term = F.normalize(torch.einsum("ijk,ik->ij", Q, zL), dim=-1)
+
+        agg_term = scatter(
+            agg_term[src], index=dst, dim=0, dim_size=N, reduce=self.reduce
+        )
+
+        loss = ((src_term - agg_term) ** 2).sum(dim=-1).mean()
+
+        return self.lamb * loss
diff --git a/topobench/nn/backbones/graph/gauge.py b/topobench/nn/backbones/graph/gauge.py
new file mode 100644
index 000000000..a740ca7f6
--- /dev/null
+++ b/topobench/nn/backbones/graph/gauge.py
@@ -0,0 +1,952 @@
+"""Riemannian Graph Foundation Model with neural vector bundles.
+
+This module implements the gauge-equivariant graph model described in
+"Are Common Substructures Transferable: Riemannian Graph Foundation Model
+with Neural Vector Bundles".
+
+The model learns, for every node, a local orthonormal frame (a gauge) that
+spans an ``r``-dimensional subspace of the ``d``-dimensional embedding space,
+smooths those frames across the graph, and uses them to update the node
+features. The building blocks correspond directly to the equations of the
+paper:
+
+- :class:`LocalCoordinatesLayer` -- equations (2)-(4).
+- :class:`GatedFlatteningLayer` -- equations (5)-(8).
+- :class:`NodeUpdateLayer` -- equations (9)-(10).
+- :class:`GaugeLayer` -- one full message-passing step combining the above.
+- :class:`GaugeModel` -- the full stack of gauge layers.
+"""
+
+import math
+from collections.abc import Callable
+
+import torch
+from torch import Tensor, nn
+from torch_geometric.utils import remove_self_loops
+from torch_scatter import scatter_add, scatter_mean, scatter_softmax
+
+activation_dict: dict[str, Callable] = {
+    "relu": nn.ReLU,
+    "leaky_relu": nn.LeakyReLU,
+    "gelu": nn.GELU,
+    "sigmoid": nn.Sigmoid,
+}
+
+
+class MultiHeadLinear(nn.Module):
+    """Per-head (per-subspace) linear layer.
+
+    Applies ``r_dim`` independent linear maps, one per head, so that head ``h``
+    transforms its own slice of the input with its own weight matrix and bias.
+    The per-head weights are stored stacked in a single parameter and applied
+    with a batched ``einsum``.
+
+    Parameters
+    ----------
+    in_channels : int
+        Number of input features per head.
+    out_channels : int
+        Number of output features per head.
+    r_dim : int
+        Number of heads (subspaces) ``r``.
+    bias : bool, optional
+        Whether each head uses a bias term (default: True).
+    device : torch.device or str or None, optional
+        Device on which to allocate the parameters (default: None).
+    dtype : torch.dtype or None, optional
+        Data type of the parameters (default: None).
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        r_dim: int,
+        bias: bool = True,
+        device=None,
+        dtype=None,
+    ):
+        super().__init__()
+
+        factory_kwargs = {"device": device, "dtype": dtype}
+
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.bias = bias
+        self.r_dim = r_dim
+
+        self.superW = nn.Parameter(
+            torch.empty(
+                self.r_dim,
+                self.out_channels,
+                self.in_channels,
+                **factory_kwargs,
+            )
+        )
+
+        if self.bias:
+            self.superB = nn.Parameter(
+                torch.empty(self.r_dim, self.out_channels, **factory_kwargs),
+            )
+        else:
+            self.register_parameter("superB", None)
+
+        self.reset_parameters()
+
+    def reset_parameters(self) -> None:
+        """Initialize the per-head weights and biases.
+
+        Uses the same scheme as :class:`torch.nn.Linear` (a uniform
+        distribution bounded by ``1 / sqrt(in_channels)``), with the fan-in
+        taken per head rather than over the stacked parameter.
+        """
+        bound = 1 / math.sqrt(self.in_channels)
+
+        torch.nn.init.uniform_(self.superW, -bound, bound)
+
+        if self.bias:
+            torch.nn.init.uniform_(self.superB, -bound, bound)
+
+    def forward(self, Z: Tensor) -> Tensor:
+        """Apply the per-head linear maps.
+
+        Parameters
+        ----------
+        Z : Tensor
+            Input tensor of shape ``[N, r, in_channels]`` where ``r`` is the
+            number of heads.
+
+        Returns
+        -------
+        Tensor
+            Output tensor of shape ``[N, r, out_channels]``.
+        """
+        if self.bias:
+            return torch.einsum("roi,Nri->Nro", self.superW, Z) + self.superB
+
+        return torch.einsum("roi,Nri->Nro", self.superW, Z)
+
+
+class MultiHeadFF(nn.Module):
+    """Per-head feed-forward network (a distinct MLP per subspace).
+
+    Stacks :class:`MultiHeadLinear` layers interleaved with activations and
+    dropout so that each of the ``r`` heads is transformed by its own
+    multi-layer perceptron. Activations and dropout are applied only between
+    layers, never after the output layer.
+
+    Parameters
+    ----------
+    in_channels : int
+        Number of input features per head.
+    out_channels : int
+        Number of output features per head.
+    r : int
+        Number of heads (subspaces) ``r``.
+    hidden_dims : list of int or None, optional
+        Widths of the hidden layers. If ``None`` the network is a single
+        per-head linear map (default: None).
+    act : str, optional
+        Name of the activation applied between layers, resolved via
+        ``activation_dict`` (default: "leaky_relu").
+    drop : float, optional
+        Dropout probability applied between layers (default: 0.0).
+    bias : bool, optional
+        Whether each per-head linear layer uses a bias term (default: True).
+    final_activation : bool, optional
+        Whether to apply the activation after the output layer as well. When
+        ``False`` the activation is applied only between layers (default: False).
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        r: int,
+        hidden_dims: list[int] | None = None,
+        act: str = "leaky_relu",
+        drop: float = 0.0,
+        bias: bool = True,
+        final_activation: bool = False,
+    ):
+        super().__init__()
+
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.act = act
+        self.dropout = drop
+        self.r = r
+        self.bias = bias
+        self.final_activation = final_activation
+
+        self.layer_sizes = [self.in_channels]
+
+        if hidden_dims is not None:
+            self.layer_sizes += [j for j in hidden_dims]
+        self.layer_sizes += [self.out_channels]
+
+        els = []
+        for j in range(len(self.layer_sizes) - 1):
+            els.append(
+                MultiHeadLinear(
+                    self.layer_sizes[j],
+                    self.layer_sizes[j + 1],
+                    r_dim=self.r,
+                    bias=self.bias,
+                )
+            )
+
+            if j < len(self.layer_sizes) - 2:
+                els.append(activation_dict[self.act]())
+                els.append(nn.Dropout(self.dropout))
+
+        # Optionally activate the output as well (e.g. the reference score_lin
+        # applies its activation to the final score before the softmax).
+        if self.final_activation:
+            els.append(activation_dict[self.act]())
+
+        self.model = nn.Sequential(*els)
+
+    def forward(self, Z: Tensor) -> Tensor:
+        """Apply the per-head feed-forward network.
+
+        Parameters
+        ----------
+        Z : Tensor
+            Input tensor of shape ``[N, r, in_channels]`` where ``r`` is the
+            number of heads.
+
+        Returns
+        -------
+        Tensor
+            Output tensor of shape ``[N, r, out_channels]``.
+        """
+        Zhat = self.model(Z)
+
+        return Zhat
+
+
+class FFBlock(nn.Module):
+    """Feed-forward block with activations, dropout and a final LayerNorm.
+
+    The block consists of ``n_hidden_layers`` hidden linear layers with an
+    activation resolved via ``activation_dict`` followed by an output linear
+    layer, with dropout applied after every layer and layer normalization
+    applied to the output.
+
+    Parameters
+    ----------
+    in_channels : int
+        Number of input features.
+    out_channels : int
+        Number of output features.
+    hidden_dim : int
+        Number of hidden units in the intermediate layers.
+    n_hidden_layers : int, optional
+        Number of hidden layers (must be at least 0). With ``0`` the block
+        reduces to a single ``Linear(in_channels, out_channels)`` (default: 1).
+    drop : float, optional
+        Dropout probability (default: 0.3).
+    bias : bool, optional
+        Whether the linear layers use a bias term (default: True).
+    act : str, optional
+        Name of the activation applied after each hidden layer, resolved via
+        ``activation_dict`` (default: "gelu").
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        hidden_dim: int,
+        n_hidden_layers: int = 1,
+        drop: float = 0.3,
+        bias: bool = True,
+        act: str = "gelu",
+    ):
+        super().__init__()
+
+        self.dropout = drop
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.hidden_dimension = hidden_dim
+        self.bias = bias
+        self.n_hidden_layers = n_hidden_layers
+        self.act = act
+
+        assert self.n_hidden_layers >= 0
+
+        els = []
+
+        # Hidden layers (none when n_hidden_layers == 0).
+        for layer_index in range(self.n_hidden_layers):
+            in_dim = (
+                self.in_channels if layer_index == 0 else self.hidden_dimension
+            )
+            els.append(
+                nn.Linear(in_dim, self.hidden_dimension, bias=self.bias)
+            )
+            els.append(activation_dict[self.act]())
+            els.append(nn.Dropout(self.dropout))
+
+        # Output layer. With no hidden layers this is a single linear map from
+        # in_channels to out_channels.
+        out_in_dim = (
+            self.in_channels
+            if self.n_hidden_layers == 0
+            else self.hidden_dimension
+        )
+        els.append(nn.Linear(out_in_dim, self.out_channels, bias=self.bias))
+        els.append(nn.Dropout(self.dropout))
+
+        self.model = torch.nn.Sequential(*els)
+        self.norm = nn.LayerNorm(self.in_channels)
+
+    def forward(self, x: Tensor) -> Tensor:
+        """Forward pass.
+
+        Parameters
+        ----------
+        x : Tensor
+            Input tensor of shape ``[..., in_channels]``.
+
+        Returns
+        -------
+        Tensor
+            Output tensor of shape ``[..., out_channels]``.
+        """
+
+        x = self.norm(x)
+        x = self.model(x)
+
+        return x
+
+
+class LocalCoordinatesLayer(torch.nn.Module):
+    """Local coordinate frame layer (equations (2)-(4)).
+
+    For each node this layer projects the node embeddings into ``r`` different
+    subspaces, aggregates a smoothed reconstruction over the neighborhood using
+    attention weights, and applies a QR decomposition to obtain, per node, an
+    orthonormal basis (a local gauge) spanning an ``r``-dimensional subspace of
+    the embedding space.
+
+    Parameters
+    ----------
+    r_subspaces : int
+        Number of subspaces (frame vectors) ``r`` learned per node.
+    d_embedd : int
+        Dimension ``d`` of the node embeddings.
+    tau : float, optional
+        Temperature used to scale the attention logits (default: 1.0).
+    bias : bool, optional
+        Whether the linear layers use a bias term (default: True).
+    act : str, optional
+        Name of the activation used by the feed-forward block, resolved via
+        ``activation_dict`` (default: "gelu").
+    f_sim_act : str, optional
+        Name of the activation used by the per-head similarity network
+        ``f_sim``, resolved via ``activation_dict`` (default: "leaky_relu").
+    dropout : float, optional
+        Dropout probability used by the feed-forward block (default: 0.3).
+    f_sim_dropout : float, optional
+        Dropout probability used by the per-head similarity network ``f_sim``
+        (default: 0.0).
+    """
+
+    # Equations (2)-(4).
+    def __init__(
+        self,
+        r_subspaces: int,
+        d_embedd: int,
+        tau: float = 1.0,
+        bias: bool = True,
+        act: str = "gelu",
+        f_sim_act: str = "leaky_relu",
+        dropout: float = 0.3,
+        f_sim_dropout: float = 0.0,
+    ):
+        super().__init__()
+
+        # A d-dimensional space admits at most d orthonormal frame vectors, so
+        # r > d would be silently truncated to d by the reduced-mode QR in
+        # forward. Reject it explicitly instead of dropping subspaces.
+        if r_subspaces > d_embedd:
+            raise ValueError(
+                f"r_subspaces ({r_subspaces}) cannot exceed d_embedd "
+                f"({d_embedd}): a d-dimensional space admits at most d "
+                "orthonormal frame vectors."
+            )
+
+        self.r = r_subspaces
+        self.tau = tau
+        self.d = d_embedd
+        self.bias = bias
+        self.act = act
+        self.f_sim_act = f_sim_act
+        self.dropout = dropout
+        self.f_sim_dropout = f_sim_dropout
+
+        # Combine the per-subspace projectors into a single nn.Linear layer;
+        # reshape into r separate projectors afterwards.
+        self.initial_projector = torch.nn.Linear(
+            self.d, self.d * self.r, bias=self.bias
+        )
+
+        # f_sim is the learnable function f of equation (3) that scores the
+        # similarity of neighboring node features (one score per subspace).
+        #
+        # Divergence from the reference implementation: the reference scores
+        # each edge with a single linear map on the raw embeddings
+        # (score_lin = Sequential(Linear(2*d -> r), LeakyReLU())). Here we use a
+        # more expressive per-subspace MLP applied to the projected multi-path
+        # embeddings Zh (shape [E, r, 2*d]), yielding the same [E, r] scores.
+        # As in the reference, final_activation applies f_sim_act (LeakyReLU by
+        # default) to the score itself before the softmax.
+        self.f_sim = MultiHeadFF(
+            2 * self.d,
+            1,
+            r=self.r,
+            hidden_dims=None,
+            act=self.f_sim_act,
+            drop=self.f_sim_dropout,
+            final_activation=True,
+        )
+
+        self.fflayer = FFBlock(
+            self.d,
+            self.d,
+            self.d,
+            bias=self.bias,
+            act=self.act,
+            drop=self.dropout,
+        )
+        self.preqr_norm = nn.LayerNorm(self.d)
+
+    def forward(self, Z: Tensor, edge_index: Tensor) -> Tensor:
+        """Forward pass computing per-node local orthonormal frames.
+
+        Implements equations (2)-(4): the neighborhood-smoothed reconstruction
+        of the node embeddings (equations (2)-(3)) followed by a QR
+        decomposition yielding an orthonormal basis per node (equation (4)).
+
+        Parameters
+        ----------
+        Z : Tensor
+            Node embeddings of shape ``[N, d]``.
+        edge_index : Tensor
+            Edge index tensor of shape ``[2, E]`` with source and destination
+            node indices.
+
+        Returns
+        -------
+        Tensor
+            Per-node orthonormal frames of shape ``[N, r, d]``.
+        """
+
+        N = Z.size(0)  # num_nodes
+        src, dst = edge_index[0], edge_index[1]
+
+        # Project the node embeddings into r different subspaces.
+        Zh = self.initial_projector(Z)  # [N, r*d]
+        Zh = Zh.reshape(N, self.r, self.d)  # [N, r, d]
+
+        # Equation (3): score each edge per subspace; f_vals has shape [E, r, 1].
+        # f_sim ends on its activation (final_activation=True), matching the
+        # reference score_lin (Linear -> LeakyReLU) before the softmax.
+        f_vals = (
+            self.f_sim(torch.concat((Zh[src], Zh[dst]), dim=-1)) / self.tau
+        )
+
+        f_vals = f_vals.squeeze(-1)  # drop the trailing singleton -> [E, r]
+        alphas = torch.softmax(f_vals, dim=-1).unsqueeze(-1)  # [E, r, 1]
+
+        # Equation (2): aggregate the weighted neighbor projections per node.
+        out = scatter_add(
+            alphas * Zh[src, :, :], index=dst, dim=0, dim_size=N
+        )  # [N, r, d]
+
+        # Normalize by the aggregated weights; clamp so that degree-0 nodes
+        # (whose scatter_add is 0) do not produce a division by zero. [N, r, 1]
+        norm = 1 / (scatter_add(alphas, dst, dim=0, dim_size=N).clamp(1e-6))
+
+        # norm * out broadcasts norm along the last dimension (d) to give
+        # [N, r, d]. Z is [N, d], so we insert a subspace axis at -2 before
+        # subtracting.
+        qhat = Z.unsqueeze(-2) - norm * out
+
+        # Feed-forward followed by a LayerNorm, then QR (equation (4)).
+        qhat = self.fflayer(qhat)
+        qhat = self.preqr_norm(qhat)
+
+        # Equation (4): qhat is [N, r, d]; transpose to [N, d, r] and apply QR
+        # to obtain an orthonormal basis per node.
+        Q, _ = torch.linalg.qr(qhat.mT)
+
+        return Q.mT
+
+
+class GatedFlatteningLayer(nn.Module):
+    """Gated flattening layer that smooths local frames (equations (5)-(8)).
+
+    This layer aligns each node's local frame with those of its neighbors. It
+    computes gating weights from the overlap between neighboring frames
+    (equation (6)), forms a gated aggregate that is blended with the original
+    frame (equation (7)), and re-orthonormalizes the result via a QR
+    decomposition (equation (8)).
+
+    Parameters
+    ----------
+    r : int
+        Number of subspaces (frame vectors) ``r`` per node.
+    gamma : float, optional
+        Blending coefficient between the original frame and the aggregated
+        neighbor frames (default: 0.01).
+    tau : float, optional
+        Temperature used to scale the gating logits (default: 1.0).
+    """
+
+    # Equations (5)-(8).
+    def __init__(self, r: int, gamma: float = 0.01, tau: float = 1.0):
+        super().__init__()
+
+        self.r = r
+        self.gamma = gamma
+        self.tau = tau
+
+    def forward(self, Q: Tensor, edge_index: Tensor) -> Tensor:
+        """Forward pass smoothing the per-node frames over the graph.
+
+        Implements equations (6)-(8): gating weights from neighboring frame
+        overlaps (equation (6)), a gated aggregate blended with the input frame
+        (equation (7)), and a final QR re-orthonormalization (equation (8)).
+
+        Parameters
+        ----------
+        Q : Tensor
+            Per-node orthonormal frames of shape ``[N, r, d]``.
+        edge_index : Tensor
+            Edge index tensor of shape ``[2, E]`` with source and destination
+            node indices.
+
+        Returns
+        -------
+        Tensor
+            Smoothed per-node orthonormal frames of shape ``[N, r, d]``.
+        """
+        # Equations (6)-(8).
+        N = Q.size(0)  # num_nodes
+        src, dst = edge_index[0], edge_index[1]
+        k = Q.size(-2)  # with k fixed, the trace of eye(k) = k
+
+        # Equation (6): gating weight from the overlap of neighboring frames.
+        # The trace of the k-by-k identity equals k.
+        g_vec = scatter_softmax(
+            ((Q[src] * Q[dst]).sum((-2, -1)) - k) / self.tau,
+            index=dst,
+            dim_size=N,
+        )
+
+        # Technical note: in principle g_ij is defined for all pairs of nodes,
+        # but because we only sum over neighbors, non-neighbor entries never
+        # contribute and need not be computed.
+
+        # Equation (7): blend the original frame with the gated neighbor
+        # aggregate.
+        Qagg = scatter_add(
+            g_vec[:, None, None] * Q[src], dim=0, index=dst, dim_size=N
+        )
+        Qhat = (1 - self.gamma) * Q + self.gamma * Qagg
+
+        # Equation (8): re-orthonormalize the blended frame with another QR.
+        Qnew, _ = torch.linalg.qr(Qhat.mT)
+
+        return Qnew.mT
+
+
+class NodeUpdateLayer(torch.nn.Module):
+    """Node feature update layer using the local frames (equations (9)-(10)).
+
+    Each node embedding is projected onto the subspace spanned by its local
+    frame and mapped through a learnable matrix (equation (9)). The projected
+    features are then aggregated over the neighborhood and, when the residual is
+    enabled, combined with a learnable transformation ``phi`` of the original
+    embedding (equation (10)). Setting ``phi_hidden_layers`` to ``None`` disables
+    the residual and recovers the reference behavior.
+
+    Parameters
+    ----------
+    in_channels : int
+        Number of input features.
+    out_channels : int
+        Number of output features.
+    phi_hidden_layers : int or None, optional
+        Number of hidden layers of the MLP residual ``phi``. With ``0`` the
+        residual is a single linear map (no hidden layer). If ``None`` the
+        residual is disabled entirely (matching the reference implementation)
+        (default: 1).
+    phi_hidden_dim : int or None, optional
+        Hidden width of the residual MLP ``phi`` (unused when
+        ``phi_hidden_layers`` is ``0``). Defaults to
+        ``max(in_channels, out_channels)`` when ``None`` (default: None).
+    act : str, optional
+        Name of the activation used by the residual MLP ``phi``, resolved via
+        ``activation_dict`` (default: "gelu").
+    dropout : float, optional
+        Dropout probability used by the residual MLP ``phi`` (default: 0.3).
+    """
+
+    # Equations (9)-(10).
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        phi_hidden_layers: int | None = 1,
+        phi_hidden_dim: int | None = None,
+        act: str = "gelu",
+        dropout: float = 0.3,
+    ):
+        super().__init__()
+
+        self.phi = None
+        # Learnable function applied to z (only when phi_hidden_layers is not None).
+        if phi_hidden_layers is not None:
+            self.phi = FFBlock(
+                in_channels=in_channels,
+                out_channels=out_channels,
+                hidden_dim=phi_hidden_dim
+                if phi_hidden_dim is not None
+                else max(in_channels, out_channels),
+                n_hidden_layers=phi_hidden_layers,
+                act=act,
+                drop=dropout,
+            )
+
+        # Learnable matrix applied to the frame-projected embedding tilde(z).
+        self.W = torch.nn.Linear(in_channels, out_channels, bias=False)
+
+    def forward(self, Z: Tensor, Q: Tensor, edge_index: Tensor) -> Tensor:
+        """Forward pass updating the node features.
+
+        Implements equations (9)-(10): the frame projection of the node
+        embeddings (equation (9)) followed by the neighborhood aggregation and,
+        when enabled, a learnable residual connection (equation (10)).
+
+        Parameters
+        ----------
+        Z : Tensor
+            Node embeddings of shape ``[N, d]``.
+        Q : Tensor
+            Per-node orthonormal frames of shape ``[N, r, d]``.
+        edge_index : Tensor
+            Edge index tensor of shape ``[2, E]`` with source and destination
+            node indices.
+
+        Returns
+        -------
+        Tensor
+            Updated node embeddings of shape ``[N, out_channels]``.
+        """
+
+        # Step 0: bind commonly used values to local names.
+        src, dst = edge_index[0], edge_index[1]
+        N = Z.size(0)  # num_nodes, for the scatter ops
+
+        # Step 1: compute the frame-projected embedding tilde(z).
+        # Equation (9): Q is [N, r, d] and Z is [N, d]; project each node
+        # embedding onto its own frame, batching over the node dimension.
+        QtZ = torch.einsum("ijk,ik->ij", Q, Z)
+        Z_tilde = torch.einsum("ikj, ik->ij", Q, QtZ)
+        Z_tilde = self.W(Z_tilde)
+
+        # Equation (10): aggregate the projected embeddings over the neighborhood.
+        #
+        # Divergence from the reference implementation: we optionally add a
+        # residual connection through self.phi. It is enabled by default and can
+        # be disabled (recovering the reference behavior) by passing
+        # phi_hidden_layers=None, in which case self.phi is None.
+        Znew = scatter_mean(Z_tilde[src], index=dst, dim=0, dim_size=N)
+
+        if self.phi is not None:
+            Znew = Znew + self.phi(Z)
+
+        return Znew
+
+
+class GaugeLayer(torch.nn.Module):
+    """A single gauge message-passing layer.
+
+    One layer computes per-node local frames with a
+    :class:`LocalCoordinatesLayer` (equations (2)-(4)), smooths them through a
+    stack of ``n_gated`` :class:`GatedFlatteningLayer` modules (equations
+    (5)-(8)), and updates the node features with a :class:`NodeUpdateLayer`
+    (equations (9)-(10)).
+
+    Parameters
+    ----------
+    d_embedd : int
+        Dimension ``d`` of the node embeddings.
+    r : int
+        Number of subspaces (frame vectors) ``r`` per node.
+    n_gated : int, optional
+        Number of gated flattening layers applied to the frames (default: 1).
+    gamma : float, optional
+        Blending coefficient used by the gated flattening layers (default: 0.01).
+    tau : float, optional
+        Temperature used to scale the attention and gating logits (default: 1.0).
+    bias : bool, optional
+        Whether the linear layers use a bias term (default: True).
+    phi_hidden_layers : int or None, optional
+        Number of hidden layers of the MLP residual ``phi`` in the node update.
+        With ``0`` the residual is a single linear map (no hidden layer). If
+        ``None`` the residual is disabled entirely (matching the reference
+        implementation) (default: 1).
+    phi_hidden_dim : int or None, optional
+        Hidden width of the residual MLP ``phi``. Defaults to ``d_embedd`` when
+        ``None`` (default: None).
+    act : str, optional
+        Name of the activation used by the feed-forward blocks, resolved via
+        ``activation_dict`` (default: "gelu").
+    f_sim_act : str, optional
+        Name of the activation used by the per-head similarity network
+        ``f_sim``, resolved via ``activation_dict`` (default: "leaky_relu").
+    dropout : float, optional
+        Dropout probability used by the feed-forward blocks (default: 0.3).
+    f_sim_dropout : float, optional
+        Dropout probability used by the per-head similarity network ``f_sim``
+        (default: 0.0).
+    """
+
+    def __init__(
+        self,
+        d_embedd: int,
+        r: int,
+        n_gated: int = 1,
+        gamma: float = 0.01,
+        tau: float = 1.0,
+        bias=True,
+        phi_hidden_layers: int | None = 1,
+        phi_hidden_dim: int | None = None,
+        act: str = "gelu",
+        f_sim_act: str = "leaky_relu",
+        dropout: float = 0.3,
+        f_sim_dropout: float = 0.0,
+    ):
+        super().__init__()
+
+        self.r = r
+        self.bias = bias
+        self.tau = tau
+        self.n_gated = n_gated
+        self.gamma = gamma
+        self.d_embedd = d_embedd
+        self.act = act
+        self.f_sim_act = f_sim_act
+        self.dropout = dropout
+        self.f_sim_dropout = f_sim_dropout
+
+        self.local_coords_layer = LocalCoordinatesLayer(
+            r_subspaces=r,
+            d_embedd=d_embedd,
+            tau=tau,
+            bias=bias,
+            act=act,
+            f_sim_act=f_sim_act,
+            dropout=dropout,
+            f_sim_dropout=f_sim_dropout,
+        )
+
+        self.gated_flattening_layers = nn.ModuleList(
+            [
+                GatedFlatteningLayer(self.r, self.gamma, self.tau)
+                for _ in range(n_gated)
+            ]
+        )
+
+        self.node_update_layer = NodeUpdateLayer(
+            self.d_embedd,
+            self.d_embedd,
+            phi_hidden_layers=phi_hidden_layers,
+            phi_hidden_dim=phi_hidden_dim,
+            act=act,
+            dropout=dropout,
+        )
+
+    def forward(self, x: Tensor, edge_index: Tensor) -> tuple[Tensor, Tensor]:
+        """Forward pass of a single gauge layer.
+
+        Parameters
+        ----------
+        x : Tensor
+            Node embeddings of shape ``[N, d]``.
+        edge_index : Tensor
+            Edge index tensor of shape ``[2, E]`` with source and destination
+            node indices.
+
+        Returns
+        -------
+        Znew : Tensor
+            Updated node embeddings of shape ``[N, d]``.
+        Q : Tensor
+            Per-node orthonormal frames of shape ``[N, r, d]``.
+        """
+
+        Q = self.local_coords_layer(x, edge_index)
+
+        for layer in self.gated_flattening_layers:
+            Q = layer(Q, edge_index)
+
+        Znew = self.node_update_layer(x, Q, edge_index)
+
+        return Znew, Q
+
+
+class GaugeModel(nn.Module):
+    """Riemannian graph foundation model with neural vector bundles.
+
+    The model first projects the input features into a ``d_embedd``-dimensional
+    embedding space and then applies a stack of :class:`GaugeLayer` modules,
+    each of which learns per-node local frames and uses them to update the node
+    features.
+
+    Parameters
+    ----------
+    n_layers : int
+        Number of gauge layers.
+    in_channels : int
+        Number of input features.
+    r : int
+        Number of subspaces (frame vectors) ``r`` per node.
+    d_embedd : int
+        Dimension ``d`` of the node embeddings.
+    n_gated : int, optional
+        Number of gated flattening layers per gauge layer (default: 1).
+    gamma : float, optional
+        Blending coefficient used by the gated flattening layers (default: 0.01).
+    tau : float, optional
+        Temperature used to scale the attention and gating logits (default: 1.0).
+    bias : bool, optional
+        Whether the linear layers use a bias term (default: True).
+    phi_hidden_layers : int or None, optional
+        Number of hidden layers of the MLP residual ``phi`` in the node update.
+        With ``0`` the residual is a single linear map (no hidden layer). If
+        ``None`` the residual is disabled entirely (matching the reference
+        implementation) (default: 1).
+    phi_hidden_dim : int or None, optional
+        Hidden width of the residual MLP ``phi``. Defaults to ``d_embedd`` when
+        ``None`` (default: None).
+    act : str, optional
+        Name of the activation used by the feed-forward blocks, resolved via
+        ``activation_dict`` (default: "gelu").
+    f_sim_act : str, optional
+        Name of the activation used by the per-head similarity network
+        ``f_sim``, resolved via ``activation_dict`` (default: "leaky_relu").
+    dropout : float, optional
+        Dropout probability used by the feed-forward blocks (default: 0.3).
+    f_sim_dropout : float, optional
+        Dropout probability used by the per-head similarity network ``f_sim``
+        (default: 0.0).
+    **kwargs : dict
+        Additional arguments.
+    """
+
+    def __init__(
+        self,
+        n_layers: int,
+        in_channels: int,
+        r: int,
+        d_embedd: int,
+        n_gated: int = 1,
+        gamma=0.01,
+        tau: float = 1.0,
+        bias=True,
+        phi_hidden_layers: int | None = 1,
+        phi_hidden_dim: int | None = None,
+        act: str = "gelu",
+        f_sim_act: str = "leaky_relu",
+        dropout: float = 0.3,
+        f_sim_dropout: float = 0.0,
+        **kwargs,
+    ):
+        super().__init__()
+
+        self.n_layers = n_layers
+        self.gamma = gamma
+        self.tau = tau
+        self.bias = bias
+        self.in_channels = in_channels
+        self.r = r
+        self.d_embedd = d_embedd
+        self.n_gated = n_gated
+        self.act = act
+        self.f_sim_act = f_sim_act
+        self.dropout = dropout
+        self.f_sim_dropout = f_sim_dropout
+
+        self.input_projector = nn.Sequential(
+            nn.Linear(in_channels, d_embedd), nn.LayerNorm(d_embedd)
+        )
+
+        self.layers = nn.ModuleList(
+            [
+                GaugeLayer(
+                    d_embedd=self.d_embedd,
+                    r=self.r,
+                    gamma=self.gamma,
+                    tau=self.tau,
+                    n_gated=self.n_gated,
+                    bias=bias,
+                    phi_hidden_layers=phi_hidden_layers,
+                    phi_hidden_dim=phi_hidden_dim,
+                    act=self.act,
+                    f_sim_act=self.f_sim_act,
+                    dropout=self.dropout,
+                    f_sim_dropout=self.f_sim_dropout,
+                )
+                for _ in range(self.n_layers)
+            ]
+        )
+
+    def forward(
+        self, x: Tensor, edge_index: Tensor, return_initial: bool = False
+    ) -> tuple[Tensor, Tensor]:
+        """Forward pass of the full model.
+
+        Parameters
+        ----------
+        x : Tensor
+            Input node features of shape ``[N, in_channels]``.
+        edge_index : Tensor
+            Edge index tensor of shape ``[2, E]`` with source and destination
+            node indices.
+        return_initial : bool
+            Whether to return the initial projection of x (default: False).
+
+        Returns
+        -------
+        z : Tensor
+            Final node embeddings of shape ``[N, d_embedd]``.
+        Q : Tensor
+            Per-node orthonormal frames of shape ``[N, r, d_embedd]`` from the
+            last gauge layer.
+        z0 : Tensor
+            The initial projection of ``x``, of shape ``[N, d_embedd]``. Only
+            returned when ``return_initial`` is True.
+        """
+        # Self-loops would make each node its own neighbor in the per-layer
+        # aggregation, which is not intended. Removed here (and, consistently,
+        # in DirichletLoss) so both aggregate over the same neighborhoods.
+        edge_index, _ = remove_self_loops(edge_index)
+
+        z = self.input_projector(x)
+
+        if return_initial:
+            z0 = z.clone()
+
+        for layer in self.layers:
+            z, Q = layer(z, edge_index)
+
+        if return_initial:
+            return z, Q, z0
+
+        return z, Q
diff --git a/topobench/nn/wrappers/graph/gauge_wrapper.py b/topobench/nn/wrappers/graph/gauge_wrapper.py
new file mode 100644
index 000000000..93f6788b7
--- /dev/null
+++ b/topobench/nn/wrappers/graph/gauge_wrapper.py
@@ -0,0 +1,36 @@
+"""Wrapper for the Gauge model."""
+
+from topobench.nn.wrappers.base import AbstractWrapper
+
+
+class GaugeWrapper(AbstractWrapper):
+    r"""Wrapper for the Gauge model.
+
+    This wrapper defines the forward pass of the model. The Gauge backbone
+    returns the rank-0 cell embeddings together with the per-node local frames;
+    only the embeddings are propagated downstream.
+    """
+
+    def forward(self, batch):
+        r"""Forward pass for the Gauge wrapper.
+
+        Parameters
+        ----------
+        batch : torch_geometric.data.Data
+            Batch object containing the batched data.
+
+        Returns
+        -------
+        dict
+            Dictionary containing the updated model output.
+        """
+        z, Q, z0 = self.backbone(
+            batch.x_0, batch.edge_index, return_initial=True
+        )
+
+        model_out = {"labels": batch.y, "batch_0": batch.batch_0}
+        model_out["x_0"] = z
+        model_out["z_0"] = z0
+        model_out["Q"] = Q
+
+        return model_out