From 7bc84d15d91898eee102fbe98dc589a0d6f40135 Mon Sep 17 00:00:00 2001 From: Vincent Moens Date: Mon, 18 Nov 2024 15:13:00 +0000 Subject: [PATCH] [Minor,Feature] Add `prefix` arg to `timeit.todict` ghstack-source-id: f1ff685caf6e8950d02dfc44ad2c1eb496495ad1 Pull Request resolved: https://github.com/pytorch/rl/pull/2576 --- torchrl/_utils.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/torchrl/_utils.py b/torchrl/_utils.py index f2ce0cf520e..0b4dd03a636 100644 --- a/torchrl/_utils.py +++ b/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():