[Minor,Feature] Add prefix arg to timeit.todict

ghstack-source-id: f1ff685
Pull Request resolved: #2576

Author: Vincent Moens

torchrl/_utils.py

@@ -117,10 +117,15 @@ def print(prefix=None) -> str:  # noqa: T202
         return "\n".join(string)
 
     @classmethod
-    def todict(cls, percall=True):
+    def todict(cls, percall=True, prefix=None):
+        def _make_key(key):
+            if prefix:
+                return f"{prefix}/{key}"
+            return key
+
         if percall:
-            return {key: val[0] for key, val in cls._REG.items()}
-        return {key: val[1] for key, val in cls._REG.items()}
+            return {_make_key(key): val[0] for key, val in cls._REG.items()}
+        return {_make_key(key): val[1] for key, val in cls._REG.items()}
 
     @staticmethod
     def erase():