We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
[BUG] timeit profiling class does not correctly track how many times a function has been called.
## Describe the bug

In file **pytorch/rl/torchrl/_utils.py**, in the class **timeit**, starting at line 31, the code is the following

```
count = self._REG[self.name][1]
self._REG[self.name][0] = (self._REG[self.name][0] * count + t) / (count + 1)
self._REG[self.name][1] = self._REG[self.name][1] + t
self._REG[self.name][2] = count + 1
```

I understand self._REG[self.name][2] tracks the number of times a function has been called.
In that case, shouldn't the first line be changed to 

```
count = self._REG[self.name][2]
```

## Checklist

- [X] I have checked that there is no similar issue in the repo (**required**)
- [  ] I have read the [documentation](https://github.com/pytorch/rl/tree/main/docs/) (**required**)
- [  ] I have provided a minimal working example to reproduce the bug (**required**)


--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `torchrl/_utils.py`
Content:
```
1 import collections
2 import math
3 import time
4 
5 import numpy as np
6 
7 
8 class timeit:
9     """A dirty but easy to use decorator for profiling code."""
10 
11     _REG = {}
12 
13     def __init__(self, name):
14         self.name = name
15 
16     def __call__(self, fn):
17         def decorated_fn(*args, **kwargs):
18             with self:
19                 out = fn(*args, **kwargs)
20                 return out
21 
22         return decorated_fn
23 
24     def __enter__(self):
25         self.t0 = time.time()
26 
27     def __exit__(self, exc_type, exc_val, exc_tb):
28         t = time.time() - self.t0
29         self._REG.setdefault(self.name, [0.0, 0.0, 0])
30 
31         count = self._REG[self.name][1]
32         self._REG[self.name][0] = (self._REG[self.name][0] * count + t) / (count + 1)
33         self._REG[self.name][1] = self._REG[self.name][1] + t
34         self._REG[self.name][2] = count + 1
35 
36     @staticmethod
37     def print(prefix=None):
38         keys = list(timeit._REG)
39         keys.sort()
40         for name in keys:
41             strings = []
42             if prefix:
43                 strings.append(prefix)
44             strings.append(
45                 f"{name} took {timeit._REG[name][0] * 1000:4.4} msec (total = {timeit._REG[name][1]} sec)"
46             )
47             print(" -- ".join(strings))
48 
49     @staticmethod
50     def erase():
51         for k in timeit._REG:
52             timeit._REG[k] = [0.0, 0.0, 0]
53 
54 
55 def _check_for_faulty_process(processes):
56     terminate = False
57     for p in processes:
58         if not p.is_alive():
59             terminate = True
60             for _p in processes:
61                 if _p.is_alive():
62                     _p.terminate()
63         if terminate:
64             break
65     if terminate:
66         raise RuntimeError(
67             "At least one process failed. Check for more infos in the log."
68         )
69 
70 
71 def seed_generator(seed):
72     """A seed generator function.
73 
74     Given a seeding integer, generates a deterministic next seed to be used in a
75     seeding sequence.
76 
77     Args:
78         seed (int): initial seed.
79 
80     Returns: Next seed of the chain.
81 
82     """
83     max_seed_val = (
84         2 ** 32 - 1
85     )  # https://discuss.pytorch.org/t/what-is-the-max-seed-you-can-set-up/145688
86     rng = np.random.default_rng(seed)
87     seed = int.from_bytes(rng.bytes(8), "big")
88     return seed % max_seed_val
89 
90 
91 class KeyDependentDefaultDict(collections.defaultdict):
92     """A key-dependent default dict.
93 
94     Examples:
95         >>> my_dict = KeyDependentDefaultDict(lambda key: "foo_" + key)
96         >>> print(my_dict["bar"])
97         foo_bar
98     """
99 
100     def __init__(self, fun):
101         self.fun = fun
102         super().__init__()
103 
104     def __missing__(self, key):
105         value = self.fun(key)
106         self[key] = value
107         return value
108 
109 
110 def prod(sequence):
111     """General prod function, that generalised usage across math and np.
112 
113     Created for multiple python versions compatibility).
114 
115     """
116     if hasattr(math, "prod"):
117         return math.prod(sequence)
118     else:
119         return int(np.prod(sequence))
120
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/torchrl/_utils.py b/torchrl/_utils.py
--- a/torchrl/_utils.py
+++ b/torchrl/_utils.py
@@ -26,12 +26,13 @@
 
     def __exit__(self, exc_type, exc_val, exc_tb):
         t = time.time() - self.t0
-        self._REG.setdefault(self.name, [0.0, 0.0, 0])
+        val = self._REG.setdefault(self.name, [0.0, 0.0, 0])
 
-        count = self._REG[self.name][1]
-        self._REG[self.name][0] = (self._REG[self.name][0] * count + t) / (count + 1)
-        self._REG[self.name][1] = self._REG[self.name][1] + t
-        self._REG[self.name][2] = count + 1
+        count = val[2]
+        N = count + 1
+        val[0] = val[0] * (count / N) + t / N
+        val[1] += t
+        val[2] = N
 
     @staticmethod
     def print(prefix=None):

{"golden_diff": "diff --git a/torchrl/_utils.py b/torchrl/_utils.py\n--- a/torchrl/_utils.py\n+++ b/torchrl/_utils.py\n@@ -26,12 +26,13 @@\n \n     def __exit__(self, exc_type, exc_val, exc_tb):\n         t = time.time() - self.t0\n-        self._REG.setdefault(self.name, [0.0, 0.0, 0])\n+        val = self._REG.setdefault(self.name, [0.0, 0.0, 0])\n \n-        count = self._REG[self.name][1]\n-        self._REG[self.name][0] = (self._REG[self.name][0] * count + t) / (count + 1)\n-        self._REG[self.name][1] = self._REG[self.name][1] + t\n-        self._REG[self.name][2] = count + 1\n+        count = val[2]\n+        N = count + 1\n+        val[0] = val[0] * (count / N) + t / N\n+        val[1] += t\n+        val[2] = N\n \n     @staticmethod\n     def print(prefix=None):\n", "issue": "[BUG] timeit profiling class does not correctly track how many times a function has been called.\n## Describe the bug\r\n\r\nIn file **pytorch/rl/torchrl/_utils.py**, in the class **timeit**, starting at line 31, the code is the following\r\n\r\n```\r\ncount = self._REG[self.name][1]\r\nself._REG[self.name][0] = (self._REG[self.name][0] * count + t) / (count + 1)\r\nself._REG[self.name][1] = self._REG[self.name][1] + t\r\nself._REG[self.name][2] = count + 1\r\n```\r\n\r\nI understand self._REG[self.name][2] tracks the number of times a function has been called.\r\nIn that case, shouldn't the first line be changed to \r\n\r\n```\r\ncount = self._REG[self.name][2]\r\n```\r\n\r\n## Checklist\r\n\r\n- [X] I have checked that there is no similar issue in the repo (**required**)\r\n- [  ] I have read the [documentation](https://github.com/pytorch/rl/tree/main/docs/) (**required**)\r\n- [  ] I have provided a minimal working example to reproduce the bug (**required**)\r\n\n", "before_files": [{"content": "import collections\nimport math\nimport time\n\nimport numpy as np\n\n\nclass timeit:\n    \"\"\"A dirty but easy to use decorator for profiling code.\"\"\"\n\n    _REG = {}\n\n    def __init__(self, name):\n        self.name = name\n\n    def __call__(self, fn):\n        def decorated_fn(*args, **kwargs):\n            with self:\n                out = fn(*args, **kwargs)\n                return out\n\n        return decorated_fn\n\n    def __enter__(self):\n        self.t0 = time.time()\n\n    def __exit__(self, exc_type, exc_val, exc_tb):\n        t = time.time() - self.t0\n        self._REG.setdefault(self.name, [0.0, 0.0, 0])\n\n        count = self._REG[self.name][1]\n        self._REG[self.name][0] = (self._REG[self.name][0] * count + t) / (count + 1)\n        self._REG[self.name][1] = self._REG[self.name][1] + t\n        self._REG[self.name][2] = count + 1\n\n    @staticmethod\n    def print(prefix=None):\n        keys = list(timeit._REG)\n        keys.sort()\n        for name in keys:\n            strings = []\n            if prefix:\n                strings.append(prefix)\n            strings.append(\n                f\"{name} took {timeit._REG[name][0] * 1000:4.4} msec (total = {timeit._REG[name][1]} sec)\"\n            )\n            print(\" -- \".join(strings))\n\n    @staticmethod\n    def erase():\n        for k in timeit._REG:\n            timeit._REG[k] = [0.0, 0.0, 0]\n\n\ndef _check_for_faulty_process(processes):\n    terminate = False\n    for p in processes:\n        if not p.is_alive():\n            terminate = True\n            for _p in processes:\n                if _p.is_alive():\n                    _p.terminate()\n        if terminate:\n            break\n    if terminate:\n        raise RuntimeError(\n            \"At least one process failed. Check for more infos in the log.\"\n        )\n\n\ndef seed_generator(seed):\n    \"\"\"A seed generator function.\n\n    Given a seeding integer, generates a deterministic next seed to be used in a\n    seeding sequence.\n\n    Args:\n        seed (int): initial seed.\n\n    Returns: Next seed of the chain.\n\n    \"\"\"\n    max_seed_val = (\n        2 ** 32 - 1\n    )  # https://discuss.pytorch.org/t/what-is-the-max-seed-you-can-set-up/145688\n    rng = np.random.default_rng(seed)\n    seed = int.from_bytes(rng.bytes(8), \"big\")\n    return seed % max_seed_val\n\n\nclass KeyDependentDefaultDict(collections.defaultdict):\n    \"\"\"A key-dependent default dict.\n\n    Examples:\n        >>> my_dict = KeyDependentDefaultDict(lambda key: \"foo_\" + key)\n        >>> print(my_dict[\"bar\"])\n        foo_bar\n    \"\"\"\n\n    def __init__(self, fun):\n        self.fun = fun\n        super().__init__()\n\n    def __missing__(self, key):\n        value = self.fun(key)\n        self[key] = value\n        return value\n\n\ndef prod(sequence):\n    \"\"\"General prod function, that generalised usage across math and np.\n\n    Created for multiple python versions compatibility).\n\n    \"\"\"\n    if hasattr(math, \"prod\"):\n        return math.prod(sequence)\n    else:\n        return int(np.prod(sequence))\n", "path": "torchrl/_utils.py"}], "after_files": [{"content": "import collections\nimport math\nimport time\n\nimport numpy as np\n\n\nclass timeit:\n    \"\"\"A dirty but easy to use decorator for profiling code.\"\"\"\n\n    _REG = {}\n\n    def __init__(self, name):\n        self.name = name\n\n    def __call__(self, fn):\n        def decorated_fn(*args, **kwargs):\n            with self:\n                out = fn(*args, **kwargs)\n                return out\n\n        return decorated_fn\n\n    def __enter__(self):\n        self.t0 = time.time()\n\n    def __exit__(self, exc_type, exc_val, exc_tb):\n        t = time.time() - self.t0\n        val = self._REG.setdefault(self.name, [0.0, 0.0, 0])\n\n        count = val[2]\n        N = count + 1\n        val[0] = val[0] * (count / N) + t / N\n        val[1] += t\n        val[2] = N\n\n    @staticmethod\n    def print(prefix=None):\n        keys = list(timeit._REG)\n        keys.sort()\n        for name in keys:\n            strings = []\n            if prefix:\n                strings.append(prefix)\n            strings.append(\n                f\"{name} took {timeit._REG[name][0] * 1000:4.4} msec (total = {timeit._REG[name][1]} sec)\"\n            )\n            print(\" -- \".join(strings))\n\n    @staticmethod\n    def erase():\n        for k in timeit._REG:\n            timeit._REG[k] = [0.0, 0.0, 0]\n\n\ndef _check_for_faulty_process(processes):\n    terminate = False\n    for p in processes:\n        if not p.is_alive():\n            terminate = True\n            for _p in processes:\n                if _p.is_alive():\n                    _p.terminate()\n        if terminate:\n            break\n    if terminate:\n        raise RuntimeError(\n            \"At least one process failed. Check for more infos in the log.\"\n        )\n\n\ndef seed_generator(seed):\n    \"\"\"A seed generator function.\n\n    Given a seeding integer, generates a deterministic next seed to be used in a\n    seeding sequence.\n\n    Args:\n        seed (int): initial seed.\n\n    Returns: Next seed of the chain.\n\n    \"\"\"\n    max_seed_val = (\n        2 ** 32 - 1\n    )  # https://discuss.pytorch.org/t/what-is-the-max-seed-you-can-set-up/145688\n    rng = np.random.default_rng(seed)\n    seed = int.from_bytes(rng.bytes(8), \"big\")\n    return seed % max_seed_val\n\n\nclass KeyDependentDefaultDict(collections.defaultdict):\n    \"\"\"A key-dependent default dict.\n\n    Examples:\n        >>> my_dict = KeyDependentDefaultDict(lambda key: \"foo_\" + key)\n        >>> print(my_dict[\"bar\"])\n        foo_bar\n    \"\"\"\n\n    def __init__(self, fun):\n        self.fun = fun\n        super().__init__()\n\n    def __missing__(self, key):\n        value = self.fun(key)\n        self[key] = value\n        return value\n\n\ndef prod(sequence):\n    \"\"\"General prod function, that generalised usage across math and np.\n\n    Created for multiple python versions compatibility).\n\n    \"\"\"\n    if hasattr(math, \"prod\"):\n        return math.prod(sequence)\n    else:\n        return int(np.prod(sequence))\n", "path": "torchrl/_utils.py"}]}