Fix (fx/tracer): revert value tracer change (Xilinx#888)

i-colbert · Mar 1, 2024 · b8a14ea · b8a14ea
1 parent 1b8ff28
commit b8a14ea
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/brevitas/fx/value_tracer.py b/src/brevitas/fx/value_tracer.py
@@ -209,7 +209,7 @@ def retrieve_method_proxy(method_args, method_kwargs):
             proxies: List[Proxy] = []
 
             def find_proxy(a):
-                if isinstance(a, cls):
+                if isinstance(a, cls) and a not in proxies:
                     proxies.append(a)
 
             fx.node.map_aggregate(method_args, find_proxy)
@@ -228,7 +228,7 @@ def find_proxy(a):
         if torch.overrides.is_tensor_method_or_property(orig_method):
             proxy = retrieve_method_proxy(args, kwargs)
             try:
-                value = orig_method(*tracer.unpack_arg(args), **tracer.unpack_arg(kwargs))
+                value = proxy.value(*tracer.unpack_arg(args), **tracer.unpack_arg(kwargs))
             except UnsetValueException:
                 value = _UNSET
             return tracer.create_proxy(