Added erase to the timeit utility (pytorch#480)

ai4co · Sep 23, 2022 · 7f9b2a9 · 7f9b2a9
1 parent c498aac
commit 7f9b2a9
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/torchrl/__init__.py b/torchrl/__init__.py
@@ -81,6 +81,11 @@ def print(prefix=None):
             )
             print(" -- ".join(strings))
 
+    @staticmethod
+    def erase():
+        for k in timeit._REG:
+            timeit._REG[k] = [0.0, 0.0, 0]
+
 
 def _check_for_faulty_process(processes):
     terminate = False