resume the finally block

HongLouyemeng · Aug 5, 2023 · 87045d4 · 87045d4
1 parent 21c33e0
commit 87045d4
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/pypose/lietensor/lietensor.py b/pypose/lietensor/lietensor.py
@@ -1248,4 +1248,7 @@ def wrapper(*args, **kwargs):
             setattr(module, name, wrap_function(func))
         yield
     finally:
-        pass
+        for func in TO_BE_WRAPPED:
+             module, name = func.__module__, func.__name__
+             module = importlib.import_module(module)
+             setattr(module, name, func)