diff --git a/examples/notebooks/Brax_Experiments_with_PGPE.ipynb b/examples/notebooks/Brax_Experiments_with_PGPE.ipynb
index e24e240..f0d2080 100644
--- a/examples/notebooks/Brax_Experiments_with_PGPE.ipynb
+++ b/examples/notebooks/Brax_Experiments_with_PGPE.ipynb
@@ -2,7 +2,7 @@
  "cells": [
   {
    "cell_type": "markdown",
-   "id": "096c868b-ef91-42df-903f-a0a046ccad95",
+   "id": "0",
    "metadata": {},
    "source": [
     "# Solving a Brax environment using EvoTorch\n",
@@ -27,7 +27,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "c786f423-220e-4c86-8e06-05edbfd45c42",
+   "id": "1",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -42,7 +42,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "8b547977-1478-4b82-95c4-12a79769336d",
+   "id": "2",
    "metadata": {},
    "source": [
     "We now check if CUDA is available. If it is, we prepare a configuration which will tell `VecGymNE` to use a single GPU both for the population and for the fitness evaluation operations. If CUDA is not available, we will instead turn to actor-based parallelization on the CPU to boost the performance."
@@ -51,7 +51,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "d45edbf5-b0e9-43a0-869a-5146321ddd4e",
+   "id": "3",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -90,7 +90,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "f1387654-13b4-4367-9d3e-ffda6b3aaf5d",
+   "id": "4",
    "metadata": {},
    "source": [
     "We now define our policy. The policy can be expressed as a string, or as an instance or as a subclass of `torch.nn.Module`."
@@ -99,7 +99,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "4d62b5e2-6471-4c86-b028-7997426270f8",
+   "id": "5",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -162,7 +162,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "166e7f59-b727-44b4-9f60-dc7cbec9943b",
+   "id": "6",
    "metadata": {},
    "source": [
     "Below, we instantiate our `VecGymNE` problem."
@@ -171,7 +171,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "56b5f554-cdec-40a5-8e21-fa09ea53e047",
+   "id": "7",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -207,7 +207,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "bce02d7c-400c-4c22-9bb8-aa70fa4b1da2",
+   "id": "8",
    "metadata": {},
    "source": [
     "---\n",
@@ -220,7 +220,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "95417793-3835-47b1-b10a-7f36e78fa3ad",
+   "id": "9",
    "metadata": {},
    "source": [
     "Initialize a PGPE to work on the problem.\n",
@@ -233,7 +233,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "bce9f851-68aa-4e67-9dbb-2474a5ebd4cf",
+   "id": "10",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -257,7 +257,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "da60f156-6756-41a4-b261-82ee62d7f7cb",
+   "id": "11",
    "metadata": {},
    "source": [
     "We register two loggers for our PGPE instance.\n",
@@ -269,7 +269,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "91270ba2-ce78-43e7-bf01-20c94b0529c3",
+   "id": "12",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -279,7 +279,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "7b4d16c1-078c-4d7d-bd6f-e3ff28173667",
+   "id": "13",
    "metadata": {},
    "source": [
     "We are now ready to start the evolutionary search."
@@ -288,7 +288,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "0a1a84a0-47ea-4592-bd37-5e96fc8f6e54",
+   "id": "14",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -298,7 +298,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "8491f968-4f43-4df6-aac0-a09c756185da",
+   "id": "15",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -310,7 +310,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "5731e007-55b3-49ef-9285-9d9137232c9d",
+   "id": "16",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -319,7 +319,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "0efa9df1-c978-4c2a-a528-c98e761caec7",
+   "id": "17",
    "metadata": {},
    "source": [
     "Now, we receive our trained policy as a torch module."
@@ -328,7 +328,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "8ea4d211-08c2-4a59-ab84-23988646895e",
+   "id": "18",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -339,7 +339,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "e7c7581d-24ba-4f06-88e2-41bb3274cd37",
+   "id": "19",
    "metadata": {},
    "source": [
     "---\n",
@@ -352,7 +352,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "c99f0ea4-7638-4e8a-a0a2-d9f483c8c096",
+   "id": "20",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -386,7 +386,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "19d38fe1-9d6f-40ec-926d-3c4066a4b66a",
+   "id": "21",
    "metadata": {},
    "source": [
     "Below, we define a utility function named `use_policy(...)`.\n",
@@ -406,7 +406,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "e8c6554d-4bdb-49b9-8c2e-956bce6ddb8a",
+   "id": "22",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -430,7 +430,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "28c044fe-4639-411c-be78-ac09cbe5e78f",
+   "id": "23",
    "metadata": {},
    "source": [
     "We now initialize a new instance of our brax environment, and trigger the jit compilation on its `reset` and `step` methods."
@@ -439,7 +439,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "c1251e82-1d0f-4e43-a6c5-1ec4c0275dab",
+   "id": "24",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -454,7 +454,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "55229cc2-aad5-4c78-b095-010a538adb40",
+   "id": "25",
    "metadata": {},
    "source": [
     "Below we run our policy and collect the states of the episodes."
@@ -463,7 +463,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "220ec837-1d1a-401e-a144-5516bdb3e493",
+   "id": "26",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -489,7 +489,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "c9bbe53a-40f5-4bf5-a063-47f1d88032d6",
+   "id": "27",
    "metadata": {},
    "source": [
     "Length of the episode and the total reward:"
@@ -498,7 +498,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "80345e9c-0694-413d-81c2-205dfacdfb5a",
+   "id": "28",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -507,7 +507,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "8b6ec424-e6cc-453a-93fd-eb07e44c1bd6",
+   "id": "29",
    "metadata": {},
    "source": [
     "Visualization of the policy:"
@@ -516,7 +516,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "7ec60419-1ad0-4f19-bc1e-f2048577ea29",
+   "id": "30",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -531,7 +531,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "a07c70f6-2c93-43a1-b4c3-edd3f395302a",
+   "id": "31",
    "metadata": {},
    "outputs": [],
    "source": [
diff --git a/examples/notebooks/Feature_Space_Illumination_with_MAPElites.ipynb b/examples/notebooks/Feature_Space_Illumination_with_MAPElites.ipynb
index 207da17..0f25c6e 100644
--- a/examples/notebooks/Feature_Space_Illumination_with_MAPElites.ipynb
+++ b/examples/notebooks/Feature_Space_Illumination_with_MAPElites.ipynb
@@ -2,7 +2,7 @@
  "cells": [
   {
    "cell_type": "markdown",
-   "id": "3cdf1a14-3dfc-41b4-af1c-d81992e993a7",
+   "id": "0",
    "metadata": {},
    "source": [
     "# Feature Space Illumination with MAPElites\n",
@@ -35,7 +35,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "205a50a2-e36c-461d-bf8e-7e1ec7db0e58",
+   "id": "1",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -49,7 +49,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "77a64b38-e6c1-47f2-9632-f93c40e7a59f",
+   "id": "2",
    "metadata": {},
    "source": [
     "Below is the definition of the bi-objective Kursawe function. The input of this function is a tensor `x` of shape `(n, 3)`, `n` being the number of solutions. The returned tensor is of shape `(n, 2)`, in which the column with index 0 stores $f_1(\\cdot)$ and column with index 1 stores $f_2(\\cdot)$."
@@ -58,7 +58,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "88dc0163-38cf-4844-9fb8-9d34457a600f",
+   "id": "3",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -79,7 +79,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "95fbd043-7bc6-4ae4-a81a-42edb92bca85",
+   "id": "4",
    "metadata": {},
    "source": [
     "Below is a wrapper for the kursawe function. The return value of the wrapped/modified Kursawe function is a tensor of shape `(n, 3)` where the column with index 0 stores $f_1(\\cdot)+f_2(\\cdot)$, column with index 1 stores $f_1(\\cdot)$, and column with index 2 stores $f_2(\\cdot)$."
@@ -88,7 +88,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "3ccda466-b2a6-4c94-a8ff-b41ded625361",
+   "id": "5",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -100,7 +100,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "ea255bd3-1f07-45ff-83de-56e76b0ea80c",
+   "id": "6",
    "metadata": {},
    "source": [
     "We now define our optimization problem. Notice that:\n",
@@ -114,7 +114,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "07966ec8-16bd-4f8a-a4db-286a33ebc4b0",
+   "id": "7",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -132,7 +132,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "bc630cfe-f364-46a9-85d0-ebc751bfc798",
+   "id": "8",
    "metadata": {},
    "source": [
     "Below, we create a hypergrid for our feature space. In our hypergrid, the global lower bound for each feature is -20, and the global upper bound for each feature is 20. For each feature, we declare that we want 50 bins."
@@ -141,7 +141,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "ce99e9e8-2e71-47af-a803-395188bd066b",
+   "id": "9",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -157,7 +157,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "f2f9c22a-a177-4ce4-9c43-34327c3da0b5",
+   "id": "10",
    "metadata": {},
    "source": [
     "Now that we have our hypergrid, we can instantiate `MAPElites`."
@@ -166,7 +166,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "c0dcaa3c-dff1-4c97-a0f3-463b10bd4ab3",
+   "id": "11",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -182,7 +182,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "3f37d8bc-25bb-4303-b747-70bdddc9f3c9",
+   "id": "12",
    "metadata": {},
    "source": [
     "Now we run our evolutionary search."
@@ -191,7 +191,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "d6a430f6-d457-4458-965f-a1e0b95f696e",
+   "id": "13",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -203,7 +203,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "611ea225-dfaa-4a51-8a1f-470c1e7444e0",
+   "id": "14",
    "metadata": {},
    "source": [
     "Below we can see how many of the cells within our hypergrid are filled:"
@@ -212,7 +212,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "153ac411-96be-4461-b693-a75d74fa5ecd",
+   "id": "15",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -221,7 +221,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "c3a7cc34-a35e-437f-9c05-4426c6f52960",
+   "id": "16",
    "metadata": {},
    "source": [
     "Now we plot our solutions.\n",
@@ -231,7 +231,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "6bb456c9-b28e-4bd9-93b4-e0a4129f1708",
+   "id": "17",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -265,7 +265,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "534f36fd-6e4e-4351-925f-cba27fffbdae",
+   "id": "18",
    "metadata": {},
    "source": [
     "---\n",
diff --git a/examples/notebooks/Genetic_Programming.ipynb b/examples/notebooks/Genetic_Programming.ipynb
index 84ff3e4..68211a4 100644
--- a/examples/notebooks/Genetic_Programming.ipynb
+++ b/examples/notebooks/Genetic_Programming.ipynb
@@ -2,7 +2,7 @@
  "cells": [
   {
    "cell_type": "markdown",
-   "id": "76de6593-c048-4427-afda-baf92f3074f0",
+   "id": "0",
    "metadata": {},
    "source": [
     "# Genetic Programming using EvoTorch\n",
@@ -17,7 +17,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "9f7661ac-c051-46fb-aecf-a0ca40886bce",
+   "id": "1",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -35,7 +35,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "1c250864-f779-4853-99b6-328634b31e1d",
+   "id": "2",
    "metadata": {},
    "source": [
     "Below are some additional functions that we wish to use in our genetic programming example."
@@ -44,7 +44,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "96d6c026-242d-4f59-8e20-c387eef61e5c",
+   "id": "3",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -83,7 +83,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "fa77555d-a491-4914-9aa4-a3cd19329251",
+   "id": "4",
    "metadata": {},
    "source": [
     "Now, we present the definition of an `Instruction`. An `Instruction` is a callable object which has access to a read-only input memory and to a runtime stack. Depending on how it was initialized, an `Instruction` can pull its arguments from the runtime stack or from the input memory. After processing its arguments, the `Instruction` will push its result onto the runtime stack."
@@ -92,7 +92,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "a3be7eb0-c829-4fe2-b52f-1928c563e112",
+   "id": "5",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -225,7 +225,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "4aaa6a9d-dc39-45c4-bc87-d7a3e7a65d89",
+   "id": "6",
    "metadata": {},
    "source": [
     "Now we define a stack-based `Interpreter`.\n",
@@ -235,7 +235,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "bf356dd6-a989-44db-9d2d-bd04aaa3bb40",
+   "id": "7",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -351,7 +351,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "2511a44d-0eb0-40be-ad55-e545a2973d4d",
+   "id": "8",
    "metadata": {},
    "source": [
     "Above, we have defined a batched interpreter where each program works on its own input and produces its own output. However, when doing symbolic regression, the most common scheme is to have fixed batch of inputs that is to be used by each program within the program batch. To be compatible with this scheme, we now define an `InterpreterWithInputBatch`, which, upon receiving a batch of inputs and a separate batch of programs, arranges them in this manner:\n",
@@ -383,7 +383,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "fb3d7e44-55d6-417e-bb44-141b35105db2",
+   "id": "9",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -452,7 +452,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "106d95a1-4073-4829-accc-d0fc4d19c8ef",
+   "id": "10",
    "metadata": {},
    "source": [
     "Now that we have our `InterpreterWithInputBatch`, we can define a `Problem` class where the goal is to minimize this mean squared error."
@@ -461,7 +461,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "396a0291-0f59-4afe-a692-fa9e1445f607",
+   "id": "11",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -552,7 +552,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "3b37a11d-4323-4626-a5e7-489688a34e02",
+   "id": "12",
    "metadata": {},
    "source": [
     "We now define a target function (the function whose definition will be searched for by our evolutionary algorithm). In the case of our example, we are searching for this function:\n",
@@ -565,7 +565,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "ea56bb31-d0a0-4fce-94c2-811b470bcad8",
+   "id": "13",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -577,7 +577,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "f25a0379-71e8-47c7-a2c7-7af1477c5c6b",
+   "id": "14",
    "metadata": {},
    "source": [
     "Below, we produce a deterministic input set, and then, using the target function, we obtain our target outputs."
@@ -586,7 +586,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "6da2c26c-64ea-4f39-8227-8596309fa19a",
+   "id": "15",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -604,7 +604,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "aaf156e5-884b-45d8-9cc9-4f53620907ce",
+   "id": "16",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -614,7 +614,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "2a22d2c9-c890-43c2-8504-0f4c15848d41",
+   "id": "17",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -634,7 +634,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "aac368a5-21b9-4b16-be94-b27497fd111e",
+   "id": "18",
    "metadata": {},
    "source": [
     "Below is a simple mutation function which changes each symbol with a probability of 10%"
@@ -643,7 +643,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "2e175a4a-09e5-431e-9805-29e46ff8b33f",
+   "id": "19",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -659,7 +659,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "4e78fa91-da7f-4685-8595-9da4a218826a",
+   "id": "20",
    "metadata": {},
    "source": [
     "Now we instantiate our genetic algorithm."
@@ -668,7 +668,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "63c50f4e-9b59-46ed-82be-f63b6479e9b5",
+   "id": "21",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -684,7 +684,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "73ef2c8e-93ab-4988-882c-69baa08f1c07",
+   "id": "22",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -694,7 +694,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "2f7aad9e-f358-4063-aac7-428feb340969",
+   "id": "23",
    "metadata": {},
    "source": [
     "Below is the best solution encountered so far, hopefully with its evaluation result expressing a near-zero error value."
@@ -703,7 +703,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "d4119755-754f-4ce5-96e1-0d473e1a3eec",
+   "id": "24",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -713,7 +713,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "400b6589-04a7-4eb3-b1fb-1c8329758ecd",
+   "id": "25",
    "metadata": {},
    "source": [
     "The program reported above can be analyzed with the help of this instruction set:"
@@ -722,7 +722,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "629dd903-eb40-47c5-96fe-729aad3c29e6",
+   "id": "26",
    "metadata": {},
    "outputs": [],
    "source": [
diff --git a/examples/notebooks/Gym_Experiments_with_PGPE_and_CoSyNE.ipynb b/examples/notebooks/Gym_Experiments_with_PGPE_and_CoSyNE.ipynb
index 7423314..841895f 100644
--- a/examples/notebooks/Gym_Experiments_with_PGPE_and_CoSyNE.ipynb
+++ b/examples/notebooks/Gym_Experiments_with_PGPE_and_CoSyNE.ipynb
@@ -2,7 +2,7 @@
  "cells": [
   {
    "cell_type": "markdown",
-   "id": "41fc2b55",
+   "id": "0",
    "metadata": {},
    "source": [
     "## Training Policies for Gym using PGPE and CoSyNE\n",
@@ -18,7 +18,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "2307fc9b",
+   "id": "1",
    "metadata": {},
    "source": [
     "## Defining the Problem\n",
@@ -29,7 +29,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "f5a506e1",
+   "id": "2",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -60,7 +60,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "37680695",
+   "id": "3",
    "metadata": {},
    "source": [
     "Now we're ready to define the problem. Let's start with the \"LunarLanderContinuous-v2\" environment."
@@ -69,7 +69,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "053e7aa9-adac-4c5f-9978-c6241ff3d4a0",
+   "id": "4",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -86,7 +86,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "bd614e9c",
+   "id": "5",
    "metadata": {},
    "source": [
     "## Creating the searcher\n",
@@ -97,7 +97,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "8ab52e4f-f6e3-4718-a101-ac1fad4d8536",
+   "id": "6",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -123,7 +123,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "96c92a5d",
+   "id": "7",
    "metadata": {},
    "source": [
     "## Training the policy\n",
@@ -134,7 +134,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "ba9a4ab6-8b7a-4dfc-8858-a5ec5df4dd0f",
+   "id": "8",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -146,7 +146,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "e6e5c870",
+   "id": "9",
    "metadata": {},
    "source": [
     "With our agent trained, it is straight-forward to now visualize the learned behaviour. For this, we will use $\\mu$, the learned center of the search distribution, as a 'best estimate' for a good policy for the environment. "
@@ -155,7 +155,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "2bd94347-a036-458c-bd30-bdb61b60e242",
+   "id": "10",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -168,7 +168,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "a99d0f8c",
+   "id": "11",
    "metadata": {},
    "source": [
     "## Training with CoSyNE\n",
@@ -179,7 +179,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "7ee0672d",
+   "id": "12",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -189,7 +189,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "ebff9be7",
+   "id": "13",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -206,7 +206,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "61b98034",
+   "id": "14",
    "metadata": {},
    "source": [
     "Defining the algorithm configuration, we aim to keep the overall evaluations-per-generation roughly the same, so use 50 individuals per generation. Additionally, we'll keep 1 elite individual-per-generation, to encourage exploitation."
@@ -215,7 +215,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "3548640c",
+   "id": "15",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -233,7 +233,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "57bdc026",
+   "id": "16",
    "metadata": {},
    "source": [
     "Once again running for 50 generations with a `StdOutLogger` attached to output the progress:"
@@ -242,7 +242,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "cc330372",
+   "id": "17",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -252,7 +252,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "3dcb5243",
+   "id": "18",
    "metadata": {},
    "source": [
     "And once again we can visualize the learned policy. As `Cosyne` is population based, it does not maintain a 'best estimate' of a good policy. Instead, we simply take the best performing solution from the current population. "
@@ -261,7 +261,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "ea3fa1fa",
+   "id": "19",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -274,7 +274,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "7a3248cc",
+   "id": "20",
    "metadata": {},
    "source": [
     "#### References\n",
diff --git a/examples/notebooks/Minimizing_Lennard-Jones_Atom_Cluster_Potentials.ipynb b/examples/notebooks/Minimizing_Lennard-Jones_Atom_Cluster_Potentials.ipynb
index 0db5826..c858132 100644
--- a/examples/notebooks/Minimizing_Lennard-Jones_Atom_Cluster_Potentials.ipynb
+++ b/examples/notebooks/Minimizing_Lennard-Jones_Atom_Cluster_Potentials.ipynb
@@ -2,7 +2,7 @@
  "cells": [
   {
    "cell_type": "markdown",
-   "id": "61bc6b75",
+   "id": "0",
    "metadata": {},
    "source": [
     "# Minimising Lennard-Jones Atom Cluster Potentials with Evolution\n",
@@ -21,7 +21,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "e67bffc6",
+   "id": "1",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -41,7 +41,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "ca81fafb",
+   "id": "2",
    "metadata": {},
    "source": [
     "Which gives a straightforward definition of the cluster potential"
@@ -50,7 +50,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "7d79f82f",
+   "id": "3",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -71,7 +71,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "d5765cc6",
+   "id": "4",
    "metadata": {},
    "source": [
     "### Obtaining Reference Solutions\n",
@@ -81,7 +81,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "a0768f59",
+   "id": "5",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -111,7 +111,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "e675fa56",
+   "id": "6",
    "metadata": {},
    "source": [
     "Now we can plot the computed potential of each solution obtained from the data base. "
@@ -120,7 +120,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "68488c43",
+   "id": "7",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -158,7 +158,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "dfd94088",
+   "id": "8",
    "metadata": {},
    "source": [
     "### Benchmarking SNES \n",
@@ -169,7 +169,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "4069d6cd",
+   "id": "9",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -219,7 +219,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "911227f4",
+   "id": "10",
    "metadata": {},
    "source": [
     "Finally let's take a look at how we did. We should see that for most cases, particularly for smaller atom clusters, SNES was either exactly recovering or was finding a solution very close to the known global optima of atom positions."
@@ -228,7 +228,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "60e37056",
+   "id": "11",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -243,7 +243,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "2d90e6f1",
+   "id": "12",
    "metadata": {},
    "source": [
     "#### References\n",
diff --git a/examples/notebooks/Model_Predictive_Control_with_CEM/reacher_mpc.ipynb b/examples/notebooks/Model_Predictive_Control_with_CEM/reacher_mpc.ipynb
index 5f8c898..03fa426 100644
--- a/examples/notebooks/Model_Predictive_Control_with_CEM/reacher_mpc.ipynb
+++ b/examples/notebooks/Model_Predictive_Control_with_CEM/reacher_mpc.ipynb
@@ -2,7 +2,7 @@
  "cells": [
   {
    "cell_type": "markdown",
-   "id": "409c9ed7-b9c6-48a8-be3b-69f90b8d777c",
+   "id": "0",
    "metadata": {},
    "source": [
     "# Model Predictive Control (MPC) with EvoTorch\n",
@@ -42,7 +42,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "17b65a39-b8c4-4771-a05c-2ca5e944b918",
+   "id": "1",
    "metadata": {},
    "source": [
     "## Requirements\n",
@@ -56,7 +56,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "35d1a364-dbae-419e-ab06-6734a7313b9a",
+   "id": "2",
    "metadata": {},
    "source": [
     "## Initial imports\n",
@@ -67,7 +67,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "2d7a725f-53e5-4a80-ad30-dbf6acb4fad0",
+   "id": "3",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -89,7 +89,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "d7302f80-3f68-407e-917c-244a47179eff",
+   "id": "4",
    "metadata": {},
    "source": [
     "## Loading the model\n",
@@ -101,7 +101,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "4264aa13-4436-42d8-a3ee-9fe46d250525",
+   "id": "5",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -119,7 +119,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "5c199a62-84e5-4455-8452-2be1b5e7c08a",
+   "id": "6",
    "metadata": {},
    "source": [
     "## Definitions\n",
@@ -130,7 +130,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "814240d5-24f7-4870-a93e-be9eef2110f7",
+   "id": "7",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -144,7 +144,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "4909a2aa-1f07-4c6e-b757-ddcbe7496a78",
+   "id": "8",
    "metadata": {},
    "source": [
     "We now define the function $\\text{predict\\_next\\_state}(s_t, a_t)$ which, given a state $s_t$ and an action $a_t$ ($t$ being the current timestep), returns the predicted next state $\\tilde{s}_{t+1}$.\n",
@@ -155,7 +155,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "acdd001d-bfc9-4574-9781-3214adb7c8d8",
+   "id": "9",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -168,7 +168,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "513ffa76-1e48-493d-bc0b-40cf15003ec3",
+   "id": "10",
    "metadata": {},
    "source": [
     "Let us now define a _plan_ $p_t$ as a series of actions planned for future timesteps, i.e.: $p_t = (a_t, a_{t+1}, a_{t+2}, ..., a_{t+(H-1)})$ where $H$ is the horizon, determining how far into the future we are planning.\n",
@@ -181,7 +181,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "7b9d7287-d9ab-4c23-abf9-c6f49c319776",
+   "id": "11",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -202,7 +202,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "06c57198-73fb-496c-80af-929b26fe28e6",
+   "id": "12",
    "metadata": {},
    "source": [
     "So far, we have defined the tools necessary for making predictions for the given plans.\n",
@@ -225,7 +225,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "824e7a59-3d16-4b67-9d96-977a3a011d87",
+   "id": "13",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -260,7 +260,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "90063fe3-7898-431a-8653-9f23237f3c0b",
+   "id": "14",
    "metadata": {},
    "source": [
     "The following is a convenience function which tackles the optimization problem defined above using the cross entropy method (CEM). The best solution produced by CEM becomes the adopted plan. Finally, the adopted plan's first action is returned (to be sent to the simulator)."
@@ -269,7 +269,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "d5bae64f-d3f8-4c2f-8422-966149671416",
+   "id": "15",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -288,7 +288,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "234223e0-78d0-4523-a842-84c8f97e38da",
+   "id": "16",
    "metadata": {},
    "source": [
     "## Main MPC Loop\n",
@@ -300,7 +300,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "7d8499b5-3cac-4f68-80f3-a09bf6accd7d",
+   "id": "17",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -310,7 +310,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "6e1886c9-09bf-41e6-a385-175a4c29e256",
+   "id": "18",
    "metadata": {},
    "source": [
     "The following function defines the main loop of MPC for a single episode of the RL environment.\n",
@@ -320,7 +320,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "c46facfd-cefd-40a2-ba2c-ac71e64698ce",
+   "id": "19",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -348,7 +348,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "89d0b224-70e4-4de6-937a-c5fb8c699ae0",
+   "id": "20",
    "metadata": {},
    "source": [
     "Run the MPC for the specified number of episodes."
@@ -357,7 +357,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "a037bd5d-16a9-48bd-8511-1338f1dd2b16",
+   "id": "21",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -369,7 +369,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "5bca6069",
+   "id": "22",
    "metadata": {},
    "source": [
     "#### References\n",
diff --git a/examples/notebooks/Model_Predictive_Control_with_CEM/train_forward_model/reacher_train.ipynb b/examples/notebooks/Model_Predictive_Control_with_CEM/train_forward_model/reacher_train.ipynb
index 4cf414f..87c4eb9 100644
--- a/examples/notebooks/Model_Predictive_Control_with_CEM/train_forward_model/reacher_train.ipynb
+++ b/examples/notebooks/Model_Predictive_Control_with_CEM/train_forward_model/reacher_train.ipynb
@@ -2,7 +2,7 @@
  "cells": [
   {
    "cell_type": "markdown",
-   "id": "bbb05aed-5dd7-41da-a88b-cfd25681b2f3",
+   "id": "0",
    "metadata": {},
    "source": [
     "# Training a forward model for `Reacher-v4`\n",
@@ -14,7 +14,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "a39524f8-82cb-4fe4-b9c7-8626dfbec077",
+   "id": "1",
    "metadata": {},
    "source": [
     "## Requirements\n",
@@ -28,7 +28,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "d928cb7f-aaea-441b-9e97-cbccbe80c73d",
+   "id": "2",
    "metadata": {},
    "source": [
     "## Initial imports\n",
@@ -39,7 +39,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "4253d704-e125-4278-9558-2333d0e5d3c8",
+   "id": "3",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -55,7 +55,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "45cc3d96-9a0f-4944-9e71-2f88c1236cbc",
+   "id": "4",
    "metadata": {},
    "source": [
     "## Declarations\n",
@@ -66,7 +66,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "3a8c5ee4-7871-4f0c-903f-8d04708699a8",
+   "id": "5",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -75,7 +75,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "f1e849e8-782b-4170-9f77-a67a8270f612",
+   "id": "6",
    "metadata": {},
    "source": [
     "By default, we use all the available CPUs of the local computer."
@@ -84,7 +84,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "fb7321a5-9c5d-41ea-900c-35a8a56d6f47",
+   "id": "7",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -93,7 +93,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "a1ffcb40-02e3-4ae9-8757-64ddaf57537f",
+   "id": "8",
    "metadata": {},
    "source": [
     "We are going to collect data from this many episodes:"
@@ -102,7 +102,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "fe0f3018-9f07-4c81-a2ed-30931279a64e",
+   "id": "9",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -111,7 +111,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "bdfd8af1-cc7f-45c9-8d54-a2c099bdc50a",
+   "id": "10",
    "metadata": {},
    "source": [
     "## Utilities for training\n",
@@ -124,7 +124,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "176e5875-32ca-49e7-a935-d119e4898644",
+   "id": "11",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -138,7 +138,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "d8fbb94c-881b-4cc9-adca-b487209f8dc4",
+   "id": "12",
    "metadata": {},
    "source": [
     "We now define a function $\\text{collect\\_data}(\\cdot)$ which collects data from multiple episodes, number of these episodes being specified via the argument `num_episodes`.\n",
@@ -155,7 +155,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "ee136854-53e1-4ede-b705-57b922cb6195",
+   "id": "13",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -190,7 +190,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "efbfc1d5-0aad-4e20-b1fc-224498162277",
+   "id": "14",
    "metadata": {},
    "source": [
     "The function below uses multiple CPUs of the local computer to collect data in parallel."
@@ -199,7 +199,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "8a4585e7-11c4-46d8-8672-fda2bd8cb9da",
+   "id": "15",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -224,7 +224,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "3630c6ff-1106-4bc1-ac27-7342a87415b2",
+   "id": "16",
    "metadata": {},
    "source": [
     "To make the supervised learning procedure more efficient, we also introduce a normalizer.\n",
@@ -238,7 +238,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "c8f14e2a-32f0-4e64-8082-1a82eeb0cc9d",
+   "id": "17",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -251,7 +251,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "2cfa5569-cab1-4c95-be99-720e2f522e8a",
+   "id": "18",
    "metadata": {},
    "source": [
     "We are now ready to collect our data and store them.\n",
@@ -263,7 +263,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "547a710f-9ea2-4626-8e55-5a2f7b4d0366",
+   "id": "19",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -280,7 +280,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "86985ad1-43a5-4c77-868e-71c9a6c43e82",
+   "id": "20",
    "metadata": {},
    "source": [
     "Below, we collect the data and their stats, and store them in the `data` namespace."
@@ -289,7 +289,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "9568dc10-b9b5-4ab9-a08f-3350f1a7fb89",
+   "id": "21",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -302,7 +302,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "b09e23f4-4273-43b6-b8fc-2078c77cf3a7",
+   "id": "22",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -312,7 +312,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "a9265d55-d586-4a0d-80e6-44cde31f903f",
+   "id": "23",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -322,7 +322,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "5baf23e8-c680-4d47-9ec6-6960168bce41",
+   "id": "24",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -331,7 +331,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "b1c0e92e-e068-447b-a54a-01068a8e7178",
+   "id": "25",
    "metadata": {},
    "source": [
     "We declare the following architecture for our neural network:"
@@ -340,7 +340,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "19ee2304-bd94-4494-b6a2-c6efdc134ef1",
+   "id": "26",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -359,7 +359,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "6211ef53-5b87-446a-b905-11791bd95da9",
+   "id": "27",
    "metadata": {},
    "source": [
     "Declare a regression problem and set the values of the hyperparameters to be used for the training procedure:"
@@ -368,7 +368,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "b3ee3386-f86f-4b84-976b-266d6e315fff",
+   "id": "28",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -384,7 +384,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "0d8387ac-c580-4d07-a065-49c04ffca608",
+   "id": "29",
    "metadata": {},
    "source": [
     "Train the model:"
@@ -393,7 +393,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "5bd260c4-cca0-478a-99b7-51d7c0160cac",
+   "id": "30",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -402,7 +402,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "5868a5aa-e6eb-428c-908e-a3d840107897",
+   "id": "31",
    "metadata": {},
    "source": [
     "At this point, we should have a trained model.\n",
@@ -413,7 +413,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "efc0e93e-f0d2-4a3c-8489-b239223daf63",
+   "id": "32",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -438,7 +438,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "779be5a8-8c51-44cd-8664-cb506e812962",
+   "id": "33",
    "metadata": {},
    "source": [
     "To compare the predictions of our model against the actual states, we instantiate a `Reacher-v4` environment."
@@ -447,7 +447,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "ebce181d-3ccd-4d40-a38b-6b8af41c0e0c",
+   "id": "34",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -457,7 +457,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "a1f491f7-aba5-4ac3-afd4-0c374daa3485",
+   "id": "35",
    "metadata": {},
    "source": [
     "In the code below, we have a loop which feeds both the actual `Reacher-v4` environment and our trained predictor the same actions.\n",
@@ -468,7 +468,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "bc7c0f0f-d136-42a6-a3bc-60692875b298",
+   "id": "36",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -515,7 +515,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "922643c0-bd2e-44c0-ad4e-ca833c638c66",
+   "id": "37",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -525,7 +525,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "cacdcf94-35b3-486e-a208-5ad1e5d577f1",
+   "id": "38",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -536,7 +536,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "b69f61a5-8844-43ed-9e8a-9a040df0e830",
+   "id": "39",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -546,7 +546,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "c8b3595f-ea70-4832-b460-e33a57901365",
+   "id": "40",
    "metadata": {},
    "source": [
     "Below, we save our trained model.\n",
@@ -556,7 +556,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "244c0305-d34e-40ba-b282-e32d3067e6ce",
+   "id": "41",
    "metadata": {},
    "outputs": [],
    "source": [
diff --git a/examples/notebooks/Variational_Quantum_Eigensolvers_with_SNES.ipynb b/examples/notebooks/Variational_Quantum_Eigensolvers_with_SNES.ipynb
index 6a4c8f0..e23c103 100644
--- a/examples/notebooks/Variational_Quantum_Eigensolvers_with_SNES.ipynb
+++ b/examples/notebooks/Variational_Quantum_Eigensolvers_with_SNES.ipynb
@@ -2,7 +2,7 @@
  "cells": [
   {
    "cell_type": "markdown",
-   "id": "f32114e2",
+   "id": "0",
    "metadata": {},
    "source": [
     "## Training a Variational Quantum Eigensolver with evotorch and pennylane\n",
@@ -36,7 +36,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "4ff6571b",
+   "id": "1",
    "metadata": {},
    "source": [
     "### Setting up the Cost Function\n",
@@ -47,7 +47,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "484054bb",
+   "id": "2",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -59,7 +59,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "9dfeba82",
+   "id": "3",
    "metadata": {},
    "source": [
     "Next, we need to calculate the electronic Hamiltonian:"
@@ -68,7 +68,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "9b29c32f",
+   "id": "4",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -81,7 +81,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "848f666d",
+   "id": "5",
    "metadata": {},
    "source": [
     "We'll be using the default PennyLane quantum simulator"
@@ -90,7 +90,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "6985795b",
+   "id": "6",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -99,7 +99,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "915de0a5",
+   "id": "7",
    "metadata": {},
    "source": [
     "Now we're ready to set up the circuit `Q` and cost function `C`. Note that due to the simplicity of the molecule, we have only a single parameter, `p=1` to optimise!"
@@ -108,7 +108,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "4bb424f1",
+   "id": "8",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -130,7 +130,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "fab9b400",
+   "id": "9",
    "metadata": {},
    "source": [
     "### Creating a EvoTorch Problem instance\n",
@@ -141,7 +141,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "491a2266-3308-42e1-9de7-b6b8b502696b",
+   "id": "10",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -155,7 +155,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "395fb09e",
+   "id": "11",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -205,7 +205,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "56f50d33",
+   "id": "12",
    "metadata": {},
    "source": [
     "### Training the VQE with EvoTorch\n",
@@ -216,7 +216,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "e648489e",
+   "id": "13",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -226,7 +226,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "4917798c",
+   "id": "14",
    "metadata": {},
    "source": [
     "And train for 100 generations"
@@ -235,7 +235,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "7166811d",
+   "id": "15",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -247,7 +247,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "d70823ae",
+   "id": "16",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -256,7 +256,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "8d1e693a",
+   "id": "17",
    "metadata": {},
    "source": [
     "Taking a look at the mean of the searcher we have:"
@@ -265,7 +265,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "e53aa9d0",
+   "id": "18",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -274,7 +274,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "e562c0dd",
+   "id": "19",
    "metadata": {},
    "source": [
     "Where the learned mean is close to the known approximate global optima of 0.208.  "
@@ -283,7 +283,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "eedd5112",
+   "id": "20",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -292,7 +292,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "d936cb93",
+   "id": "21",
    "metadata": {},
    "source": [
     "### A more challenging example\n",
@@ -305,7 +305,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "38cb1de3",
+   "id": "22",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -315,7 +315,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "601b02f2",
+   "id": "23",
    "metadata": {},
    "source": [
     "And as before we need the electronic Hamiltonian"
@@ -324,7 +324,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "9c9649cb",
+   "id": "24",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -344,7 +344,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "4895f351",
+   "id": "25",
    "metadata": {},
    "source": [
     "Making a new cost function, this time using UCCSD ansatz"
@@ -353,7 +353,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "9e4a57b2",
+   "id": "26",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -389,7 +389,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "58015428",
+   "id": "27",
    "metadata": {},
    "source": [
     "Putting this together in a problem definition:"
@@ -398,7 +398,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "cd245923",
+   "id": "28",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -465,7 +465,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "f2bc7456",
+   "id": "29",
    "metadata": {},
    "source": [
     "And from there, all we need to do is train!"
@@ -474,7 +474,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "dfa1e634",
+   "id": "30",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -488,7 +488,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "4222cde0",
+   "id": "31",
    "metadata": {},
    "source": [
     "And visualize the progress:"
@@ -497,7 +497,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "187f38e5",
+   "id": "32",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -507,7 +507,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "e5b7906b",
+   "id": "33",
    "metadata": {},
    "source": [
     "#### References\n",