[src] Allow disabling script_if_tracing for ONNX export (#16)

asteroid-team · Jul 30, 2022 · e15bb3f · e15bb3f
1 parent 8a3d13f
commit e15bb3f
Showing 1 changed file with 15 additions and 1 deletion.
diff --git a/asteroid_filterbanks/scripting.py b/asteroid_filterbanks/scripting.py
@@ -2,6 +2,20 @@
 import torch
 
 
+global SCRIPT_ENABLED
+SCRIPT_ENABLED = True
+
+
+def disable_script_if_tracing():
+    global SCRIPT_ENABLED
+    SCRIPT_ENABLED = False
+
+
+def enable_script_if_tracing():
+    global SCRIPT_ENABLED
+    SCRIPT_ENABLED = True
+
+
 def is_tracing():
     # Taken for pytorch for compat in 1.6.0
     """
@@ -32,7 +46,7 @@ def script_if_tracing(fn):
 
     @functools.wraps(fn)
     def wrapper(*args, **kwargs):
-        if not is_tracing():
+        if not is_tracing() or not SCRIPT_ENABLED:
             # Not tracing, don't do anything
             return fn(*args, **kwargs)