dev(hansbug): fix MoductDict register

opendilab · Oct 22, 2023 · a323641 · a323641
1 parent 4f9a9f1
commit a323641
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 9 deletions.
diff --git a/treevalue/tree/integration/jax.py b/treevalue/tree/integration/jax.py
@@ -1,8 +1,6 @@
 import warnings
 from functools import wraps
 
-from ..tree import register_dict_type
-
 try:
     import jax
     from jax.tree_util import register_pytree_node
@@ -22,10 +20,3 @@ def register_for_jax(cls):
 
     register_for_jax(TreeValue)
     register_for_jax(FastTreeValue)
-
-try:
-    from torch.nn import ModuleDict
-except (ModuleNotFoundError, ImportError):
-    pass
-else:
-    register_dict_type(ModuleDict, ModuleDict.items)
diff --git a/treevalue/tree/integration/torch.py b/treevalue/tree/integration/torch.py
@@ -1,6 +1,8 @@
 import warnings
 from functools import wraps
 
+from ..tree import register_dict_type
+
 try:
     import torch
     from torch.utils._pytree import _register_pytree_node
@@ -20,3 +22,10 @@ def register_for_torch(cls):
 
     register_for_torch(TreeValue)
     register_for_torch(FastTreeValue)
+
+try:
+    from torch.nn import ModuleDict
+except (ModuleNotFoundError, ImportError):
+    pass
+else:
+    register_dict_type(ModuleDict, ModuleDict.items)