From c024bc9b5655899bff2fa3da7829ba28f4cc621b Mon Sep 17 00:00:00 2001
From: Xietao Wang Lin <Xietao.Wang-Lin@warwick.ac.uk>
Date: Fri, 8 Aug 2025 14:26:33 +0100
Subject: [PATCH] remove missing cfg.box

---
 lagrangebench/runner.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/lagrangebench/runner.py b/lagrangebench/runner.py
index 26e1ecc..292a041 100644
--- a/lagrangebench/runner.py
+++ b/lagrangebench/runner.py
@@ -244,7 +244,8 @@ def model_fn(x):
 
         MODEL = models.SEGNN
     elif model_name == "egnn":
-        box = cfg.box
+        bounds = np.array(metadata["bounds"])
+        box = bounds[:, 1] - bounds[:, 0]
         if jnp.array(metadata["periodic_boundary_conditions"]).any():
             displacement_fn, shift_fn = space.periodic(jnp.array(box))
         else: