Merge pull request #5 from opendilab/fix/walk

fix(hansbug): fix uncompitable problem with walk
opendilab · Mar 3, 2022 · d664fb5 · d664fb5
2 parents 1e922e0 + 61f5e94
commit d664fb5
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 4 deletions.
diff --git a/requirements.txt b/requirements.txt
@@ -1,4 +1,4 @@
-treevalue>=1.2.0
+treevalue>=1.2.2
 torch>=1.1.0,<=1.10.0
 hbutils>=0.0.1
 numpy
diff --git a/treetensor/common/wrappers.py b/treetensor/common/wrappers.py
@@ -1,7 +1,6 @@
 from functools import wraps
-from operator import itemgetter
 
-from treevalue import TreeValue, walk
+from treevalue import TreeValue, flatten_values
 
 __all__ = [
     'ireduce',
@@ -17,7 +16,7 @@ def _decorator(func):
         def _new_func(*args, **kwargs):
             result = func(*args, **kwargs)
             if isinstance(result, TreeValue):
-                it = map(itemgetter(1), walk(result, include_nodes=False))
+                it = flatten_values(result)
                 return rfunc(piter(it))
             else:
                 return result