diff --git a/tf_pwa/experimental/wrap_function.py b/tf_pwa/experimental/wrap_function.py
index cf2b4768..5f334f32 100644
--- a/tf_pwa/experimental/wrap_function.py
+++ b/tf_pwa/experimental/wrap_function.py
@@ -4,7 +4,9 @@
 
 def _wrap_struct(dic, first_none=True):
     if isinstance(dic, dict):
-        return {k: _wrap_struct(v, first_none) for k, v in dic.items()}
+        return {
+            k: _wrap_struct(dic[k], first_none) for k in sorted(dic.keys())
+        }
     if isinstance(dic, list):
         return [_wrap_struct(v, first_none) for v in dic]
     if isinstance(dic, tuple):
@@ -19,8 +21,8 @@ def _wrap_struct(dic, first_none=True):
 
 def _flatten(dic):
     if isinstance(dic, dict):
-        for k, v in dic.items():
-            yield from _flatten(v)
+        for k in sorted(dic.keys()):
+            yield from _flatten(dic[k])
     if isinstance(dic, (list, tuple)):
         for v in dic:
             yield from _flatten(v)
diff --git a/tf_pwa/tests/test_full.py b/tf_pwa/tests/test_full.py
index 3ff2d113..8cf34227 100644
--- a/tf_pwa/tests/test_full.py
+++ b/tf_pwa/tests/test_full.py
@@ -71,6 +71,13 @@ def toy_config(gen_toy):
     return config
 
 
+@pytest.fixture
+def toy_config_lazy(gen_toy):
+    config = ConfigLoader(f"{this_dir}/config_lazycall.yml")
+    config.set_params(f"{this_dir}/exp_params.json")
+    return config
+
+
 def test_build_angle_amplitude(toy_config):
     data = toy_config.get_data("data")
     dec = toy_config.get_amplitude().decay_group
@@ -236,6 +243,13 @@ def test_fit(toy_config, fit_result):
     xy_err = pt.get_error_matrix([x, y])
 
 
+def test_lazycall(toy_config_lazy):
+    toy_config_lazy.fit(batch=100000)
+    toy_config_lazy.plot_partial_wave(
+        prefix="toy_data/figure_lazy", batch=100000
+    )
+
+
 def test_cal_chi2(toy_config, fit_result):
     toy_config.cal_chi2(bins=[[2, 2]] * 2, mass=["R_BD", "R_CD"])