fix modules with dict (#819)

2025-09-18 18:28:12 +08:00 · 2024-03-12 08:54:06 -07:00
parent 8e5600022a
commit 366478c560
2 changed files with 13 additions and 1 deletions
--- a/python/mlx/nn/layers/base.py
+++ b/python/mlx/nn/layers/base.py
@@ -20,7 +20,7 @@ def _unwrap(model, value_key, value, filter_fn, map_fn, is_leaf_fn):

    elif isinstance(value, dict):
        nd = {}
-        for k, v in v.items():
+        for k, v in value.items():
            tk = f"{value_key}.{k}"
            nd[k] = (
                _unwrap(model, tk, v, filter_fn, map_fn, is_leaf_fn)