fix an op convertion error (#469)

Address this [todo](#458 (comment))
siliconflow · Dec 28, 2023 · 7330b8c · 7330b8c
1 parent 2d7793f
commit 7330b8c
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 9 deletions.
diff --git a/recordconda b/recordconda
diff --git a/src/onediff/infer_compiler/transform/builtin_transform.py b/src/onediff/infer_compiler/transform/builtin_transform.py
@@ -9,7 +9,6 @@
 import torch
 import oneflow as flow
 
-
 from .manager import transform_mgr
 from ..utils.log_utils import logger
 from ..utils.patch_for_diffusers import diffusers_checker
@@ -325,14 +324,15 @@ def _(mod: None, verbose=False):
 def _(mod: types.BuiltinFunctionType, verbose=False):
     if hasattr(mod, "__module__"):
         mod_name = None
-        #TODO: This solution is a compromise for now.
-        #TODO: Should register nn.linear later to solve it elegantly 
-        if mod == torch._C._nn.linear:
-            return flow.nn.functional.linear
-        elif mod.__module__.startswith("torch._C._nn"):
-            mod_name = mod.__module__.replace(
-                "torch._C._nn", "oneflow._oneflow_internal._C"
-            )
+        if mod.__module__.startswith("torch._C._nn"):
+            # The equivalence of mod inside torch._C._nn may be
+            # defined in flow.nn.functional
+            if getattr(flow.nn.functional, mod.__name__):
+                mod_name = "oneflow.nn.functional"
+            else:
+                mod_name = mod.__module__.replace(
+                    "torch._C._nn", "oneflow._oneflow_internal._C"
+                )
         elif mod.__module__.startswith("torch"):
             try:
                 if getattr(torch.nn.functional, mod.__name__) == mod: