remove ref

Signed-off-by: Wenqi Li <wenqil@nvidia.com>

tests/utils.py

@@ -440,6 +440,7 @@ def _wrapper(*args, **kwargs):
             for p in processes:
                 p.join()
                 assert results.get(), "Distributed call failed."
+            _del_original_func(obj)
 
         return _wrapper
 
@@ -521,6 +522,7 @@ def _wrapper(*args, **kwargs):
             finally:
                 p.join()
 
+            _del_original_func(obj)
             res = None
             try:
                 res = results.get(block=False)
@@ -546,6 +548,15 @@ def _cache_original_func(obj) -> None:
     _original_funcs[obj.__name__] = obj
 
 
+def _del_original_func(obj):
+    """pop the original function from cache."""
+    global _original_funcs
+    _original_funcs.pop(obj.__name__, None)
+    if torch.cuda.is_available():  # clean up the cached function
+        torch.cuda.synchronize()
+        torch.cuda.empty_cache()
+
+
 def _call_original_func(name, module, *args, **kwargs):
     if name not in _original_funcs:
         _original_module = importlib.import_module(module)  # reimport, refresh _original_funcs