Loosen golden test tolerance for CPU/GPU cross-platform compatibility

Author: cool-RR

tests/test_highjax_trainer/test_golden_runs/common.py

@@ -38,7 +38,7 @@ def check_golden_run(epoch_metrics: list[dict], golden_data: dict) -> None:
             expected = expected_values[i]
             actual = float(metrics[key])
 
-            if not jnp.isclose(actual, expected, rtol=1e-04, atol=1e-07):
+            if not jnp.isclose(actual, expected, rtol=1e-02, atol=1e-05):
                 bad_fields.append((key, actual, expected))
 
         if bad_fields: