diff --git a/examples/notebooks/Brax_Experiments_Visualization.ipynb b/examples/notebooks/Brax_Experiments_Visualization.ipynb
index b8f9b0e..4c15ea8 100644
--- a/examples/notebooks/Brax_Experiments_Visualization.ipynb
+++ b/examples/notebooks/Brax_Experiments_Visualization.ipynb
@@ -2,7 +2,7 @@
  "cells": [
   {
    "cell_type": "markdown",
-   "id": "fc7b28eb-db2b-4a02-87d4-dfc44a9fecc2",
+   "id": "0",
    "metadata": {},
    "source": [
     "# Visualization of the `brax` Experiment Results\n",
@@ -13,7 +13,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "d5350a77-c84b-4d05-9e1c-4a6e0e2f6d0d",
+   "id": "1",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -24,7 +24,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "e4c4d0df-0480-4c5f-a4ab-d34727286cc1",
+   "id": "2",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -34,7 +34,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "d64f3d04-680a-43c5-a807-dbaf9f917198",
+   "id": "3",
    "metadata": {},
    "source": [
     "---"
@@ -43,7 +43,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "7025f8b3-9982-4196-9790-9e8c1491dc06",
+   "id": "4",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -55,7 +55,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "d5bccb74-a64a-4e55-9ae9-9686f772fa60",
+   "id": "5",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -66,7 +66,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "0d71c7c3-bb36-47fb-88dc-bd285c54c44a",
+   "id": "6",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -77,7 +77,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "54c7bee4-bf5c-445c-b009-40484391461c",
+   "id": "7",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -89,7 +89,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "30eca014-be71-4162-b675-0eaf8b037bce",
+   "id": "8",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -100,7 +100,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "78ce1f2f-cb89-432e-ad98-07f7e1d5d457",
+   "id": "9",
    "metadata": {},
    "source": [
     "Below, we put the values of the center solution into the policy network as parameters:"
@@ -109,7 +109,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "545d693d-ccd9-4ea7-b068-06b832f32be5",
+   "id": "10",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -118,7 +118,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "e7c7581d-24ba-4f06-88e2-41bb3274cd37",
+   "id": "11",
    "metadata": {},
    "source": [
     "---\n",
@@ -131,7 +131,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "8d798a71-003f-4bed-9f12-391dc7b797c0",
+   "id": "12",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -143,7 +143,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "30080d4b-227c-4bde-8009-7a97c6e177f1",
+   "id": "13",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -153,7 +153,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "c64a6482-0533-40c8-b8f9-1b1ad3ee0019",
+   "id": "14",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -163,7 +163,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "ee69ba99-bcb7-4d21-9f5c-6c330cb3ed96",
+   "id": "15",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -183,7 +183,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "c99f0ea4-7638-4e8a-a0a2-d9f483c8c096",
+   "id": "16",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -194,7 +194,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "19d38fe1-9d6f-40ec-926d-3c4066a4b66a",
+   "id": "17",
    "metadata": {},
    "source": [
     "Below, we define a utility function named `use_policy(...)`.\n",
@@ -214,7 +214,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "e8c6554d-4bdb-49b9-8c2e-956bce6ddb8a",
+   "id": "18",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -238,7 +238,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "28c044fe-4639-411c-be78-ac09cbe5e78f",
+   "id": "19",
    "metadata": {},
    "source": [
     "We now initialize a new instance of our brax environment, and trigger the jit compilation on its `reset` and `step` methods."
@@ -247,7 +247,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "c1251e82-1d0f-4e43-a6c5-1ec4c0275dab",
+   "id": "20",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -258,7 +258,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "55229cc2-aad5-4c78-b095-010a538adb40",
+   "id": "21",
    "metadata": {},
    "source": [
     "Below we run our policy and collect the states of the episodes."
@@ -267,7 +267,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "220ec837-1d1a-401e-a144-5516bdb3e493",
+   "id": "22",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -291,7 +291,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "c9bbe53a-40f5-4bf5-a063-47f1d88032d6",
+   "id": "23",
    "metadata": {},
    "source": [
     "Length of the episode and the total reward:"
@@ -300,7 +300,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "80345e9c-0694-413d-81c2-205dfacdfb5a",
+   "id": "24",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -310,7 +310,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "8d2e2506-5bb2-4b1f-bb00-004465008fa3",
+   "id": "25",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -325,7 +325,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "8b6ec424-e6cc-453a-93fd-eb07e44c1bd6",
+   "id": "26",
    "metadata": {},
    "source": [
     "Visualization of the policy:"
@@ -334,7 +334,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "29b9c3aa-8068-40bc-9ec2-e5ff759ed22c",
+   "id": "27",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -349,7 +349,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "b50e36a7-dd33-4fc3-9f9d-2d7dca93110b",
+   "id": "28",
    "metadata": {},
    "outputs": [],
    "source": [
diff --git a/examples/notebooks/Brax_Experiments_with_PGPE.ipynb b/examples/notebooks/Brax_Experiments_with_PGPE.ipynb
index df0aacd..5a86360 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": [
@@ -46,7 +46,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."
@@ -55,7 +55,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "8fb0cb8e-c247-4416-94cd-9a8349192421",
+   "id": "3",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -88,7 +88,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "ef60c766-d03a-4c37-8966-1db77d7658f9",
+   "id": "4",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -99,7 +99,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "d45edbf5-b0e9-43a0-869a-5146321ddd4e",
+   "id": "5",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -156,7 +156,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "f1387654-13b4-4367-9d3e-ffda6b3aaf5d",
+   "id": "6",
    "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`."
@@ -165,7 +165,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "4d62b5e2-6471-4c86-b028-7997426270f8",
+   "id": "7",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -228,7 +228,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "166e7f59-b727-44b4-9f60-dc7cbec9943b",
+   "id": "8",
    "metadata": {},
    "source": [
     "Below, we instantiate our `VecGymNE` problem."
@@ -237,7 +237,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "94921740-188b-4b57-84a2-d5fab5b04a26",
+   "id": "9",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -246,7 +246,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "b502af4e-ff98-4483-af4b-6109ac768d36",
+   "id": "10",
    "metadata": {},
    "source": [
     "**Note.**\n",
@@ -256,7 +256,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "56b5f554-cdec-40a5-8e21-fa09ea53e047",
+   "id": "11",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -290,7 +290,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "95417793-3835-47b1-b10a-7f36e78fa3ad",
+   "id": "12",
    "metadata": {},
    "source": [
     "Initialize a PGPE to work on the problem.\n",
@@ -303,7 +303,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "bce9f851-68aa-4e67-9dbb-2474a5ebd4cf",
+   "id": "13",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -334,7 +334,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "ef3ad66e-e11c-4fa4-b250-e6786ce94270",
+   "id": "14",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -347,7 +347,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "da60f156-6756-41a4-b261-82ee62d7f7cb",
+   "id": "15",
    "metadata": {},
    "source": [
     "We register two loggers for our PGPE instance.\n",
@@ -359,7 +359,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "91270ba2-ce78-43e7-bf01-20c94b0529c3",
+   "id": "16",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -369,7 +369,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "7b4d16c1-078c-4d7d-bd6f-e3ff28173667",
+   "id": "17",
    "metadata": {},
    "source": [
     "We are now ready to start the evolutionary search."
@@ -378,7 +378,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "0a1a84a0-47ea-4592-bd37-5e96fc8f6e54",
+   "id": "18",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -392,7 +392,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "8491f968-4f43-4df6-aac0-a09c756185da",
+   "id": "19",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -403,7 +403,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "1951f337-91ed-40a3-b432-600c8380c286",
+   "id": "20",
    "metadata": {},
    "source": [
     "See the notebook [Brax_Experiments_Visualization.ipynb](Brax_Experiments_Visualization.ipynb) for visualizing the pickle files generated by this notebook."
diff --git a/examples/notebooks/Evolving_Objects.ipynb b/examples/notebooks/Evolving_Objects.ipynb
index ca2e88c..668b27f 100644
--- a/examples/notebooks/Evolving_Objects.ipynb
+++ b/examples/notebooks/Evolving_Objects.ipynb
@@ -2,7 +2,7 @@
  "cells": [
   {
    "cell_type": "markdown",
-   "id": "471df805-4945-4039-bf8a-b1fcfc50e990",
+   "id": "0",
    "metadata": {},
    "source": [
     "# Evolving Objects\n",
@@ -57,7 +57,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "b3195d2e-bdaa-41b8-aafd-45b1b0d9e6f1",
+   "id": "1",
    "metadata": {},
    "source": [
     "We begin with importing the necessary libraries:"
@@ -66,7 +66,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "e296cf9d-86c2-4952-9df3-318e04866766",
+   "id": "2",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -88,7 +88,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "a8d4cdc7-555e-48f6-be8a-cca47ac41173",
+   "id": "3",
    "metadata": {},
    "source": [
     "Below are utility functions for generating and loading observation normalization data."
@@ -97,7 +97,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "70c6cebe-435a-40ba-a6ef-6e217a701cc1",
+   "id": "4",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -122,7 +122,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "a6afff2a-e216-4141-ab99-1d90cc4d6043",
+   "id": "5",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -217,7 +217,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "ebaae4a9-4013-404a-a5af-530dbf5d53ed",
+   "id": "6",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -252,7 +252,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "0140ac43-4c60-42ee-80b6-eb69516b5625",
+   "id": "7",
    "metadata": {},
    "source": [
     "Below is the environment ID of the MuJoCo Ant environment that we are targeting.\n",
@@ -263,7 +263,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "0d29932a-c56d-4147-b1b9-db046f9f726f",
+   "id": "8",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -272,7 +272,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "a8fd4bdf-202d-4796-a45b-24b30fc15f95",
+   "id": "9",
    "metadata": {},
    "source": [
     "Below, we generate (or load) observation normalization data for the environment:"
@@ -281,7 +281,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "626d76d1-05ce-45ab-8b1a-cb6490aece1d",
+   "id": "10",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -292,7 +292,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "215d251c-62c7-4bff-957e-68447c44ce97",
+   "id": "11",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -584,7 +584,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "42e14515-0f8c-4cb8-9c37-eed22631cc34",
+   "id": "12",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -629,7 +629,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "2fe8f8e0-ca6c-467b-be17-e1428cb7e455",
+   "id": "13",
    "metadata": {},
    "source": [
     "Now that we have the implementation of our `Problem` subclass, we instantiate it:"
@@ -638,7 +638,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "fc1893a3-0bf9-4bd0-abac-bc136d8d1db2",
+   "id": "14",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -652,7 +652,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "26241767-d6cc-49ae-bede-e895dd944536",
+   "id": "15",
    "metadata": {},
    "source": [
     "Instantiate the `GeneticAlgorithm`"
@@ -661,7 +661,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "8ad83512-8019-4e10-9d91-8314b22037d5",
+   "id": "16",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -699,7 +699,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "2610bcab-8e2f-4422-bb0a-66f01a37ea35",
+   "id": "17",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -710,7 +710,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "337daf81-10f4-4d00-8532-c026d6453129",
+   "id": "18",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -723,7 +723,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "f3005822-a3a7-4bab-88e4-41f526fb78cb",
+   "id": "19",
    "metadata": {},
    "source": [
     "---\n",
@@ -734,7 +734,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "d043c619-7e7d-4ab5-a105-2f8b51c45a77",
+   "id": "20",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -755,7 +755,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "1835b151-0f6f-4028-a5e4-51838d64f659",
+   "id": "21",
    "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/Functional_API/batched_searches.ipynb b/examples/notebooks/Functional_API/batched_searches.ipynb
index 69042e5..cdfdf72 100644
--- a/examples/notebooks/Functional_API/batched_searches.ipynb
+++ b/examples/notebooks/Functional_API/batched_searches.ipynb
@@ -2,7 +2,7 @@
  "cells": [
   {
    "cell_type": "markdown",
-   "id": "d66277ff-bb08-48c6-b1b3-21919b700c9b",
+   "id": "0",
    "metadata": {},
    "source": [
     "# Maintaining a batch of populations using the functional EvoTorch API\n",
@@ -22,7 +22,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "18d7c1eb-3c93-4c05-9044-8cfd9334e5a9",
+   "id": "1",
    "metadata": {},
    "source": [
     "We begin by importing the necessary libraries and defining some useful variables:"
@@ -31,7 +31,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "c5971b63-d6b9-4c02-9669-ab2060af833c",
+   "id": "2",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -50,7 +50,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "6461d4b6-9a70-434a-8aa1-e7be410d96d3",
+   "id": "3",
    "metadata": {},
    "source": [
     "## Ask-and-tell\n",
@@ -60,7 +60,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "a804967a-1845-40c6-a80d-9424edfa9052",
+   "id": "4",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -121,7 +121,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "5e672e8e-77cf-439d-82fb-1775dbe6fc69",
+   "id": "5",
    "metadata": {},
    "source": [
     "## The @rowwise decorator\n",
@@ -132,7 +132,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "44738506-c134-46cf-8268-9f88ca358643",
+   "id": "6",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -145,7 +145,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "2a426eb1-b2ed-4694-8184-ef8bbccbf73b",
+   "id": "7",
    "metadata": {},
    "source": [
     "Notice how the fitness function above is decorated via `@rowwise`. This decorator tells EvoTorch that the user has defined the function to operate on its argument `x` as a vector (i.e. a 1-dimensional tensor). This makes it conceptually easier to implement the function and helps EvoTorch safely `vmap()` it in order to apply it to populations or a batch of populations as needed. `@rowwise` ensures that:\n",
@@ -161,7 +161,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "36be528e-fae4-4104-a27a-b255c7facab1",
+   "id": "8",
    "metadata": {},
    "source": [
     "## Batched (vectorized) searches\n",
@@ -172,7 +172,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "046dd528-ee4c-469a-96f9-a5bef4ca2bd5",
+   "id": "9",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -181,7 +181,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "b9434b49-90c7-439f-81ae-db4ad024cb91",
+   "id": "10",
    "metadata": {},
    "source": [
     "For both functional `cem` and functional `pgpe`, the hyperparameter `stdev_max_change` can be given as a scalar (which then will be expanded to a vector), or as a vector (which then will be used as it is), or as a batch of vectors (which will mean that for each batch item `i`, the `i`-th `stdev_max_change` vector will be used).\n",
@@ -192,7 +192,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "e97787a3-8ed6-4db8-8396-1cc3ae19cb59",
+   "id": "11",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -204,7 +204,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "b28047db-e09a-4327-bf6d-e12b0dd9b8f1",
+   "id": "12",
    "metadata": {},
    "source": [
     "Next we simply provide these to the CEM state initializer and execute CEM using the ask-and-tell interface exacty as before. Internally, EvoTorch will recognize the new batch dimension and appropriately `vmap()` the fitness function for us!"
@@ -213,7 +213,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "cba9629f-afc5-482f-bd36-14edb01b9de9",
+   "id": "13",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -256,7 +256,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "6c6d32a5-c4b0-4497-9a92-e82a9cfdd3c1",
+   "id": "14",
    "metadata": {},
    "source": [
     "If this notebook is executed on a GPU, the above batched search will take less time than `batch_size` times the time taken by the single search above, particularly for larger values of `batch_size`. Here are the center points found by CEM:"
@@ -265,7 +265,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "b370baa4-70c6-4093-87d0-36fa63eabe19",
+   "id": "15",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -274,7 +274,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "cc1ceb7d-4c6b-4cdc-a7e8-e52232819b61",
+   "id": "16",
    "metadata": {},
    "source": [
     "## Another example\n",
@@ -289,7 +289,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "d64a199c-bb51-468e-80dd-393dc0accc9e",
+   "id": "17",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -299,7 +299,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "0c3bcc5c-471f-49d2-a1f5-0f7ef33ccfeb",
+   "id": "18",
    "metadata": {},
    "source": [
     "Now we prepare the first state of our `pgpe` search:"
@@ -308,7 +308,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "9c3674f4-35df-406a-9574-7b1e417717f7",
+   "id": "19",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -348,7 +348,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "aad76f22-28dc-42c3-b504-34c585130c49",
+   "id": "20",
    "metadata": {},
    "source": [
     "Below is the main loop of the evolutionary search."
@@ -357,7 +357,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "5244f8bd-b3a1-49e7-9ce3-1b78ae649f3b",
+   "id": "21",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -387,7 +387,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "e614ed61-66ab-402a-890b-db376ba46ab8",
+   "id": "22",
    "metadata": {},
    "source": [
     "Here are the center points found by `pgpe`:"
@@ -396,7 +396,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "4da70d95-8828-4aac-917b-edb0a19e304f",
+   "id": "23",
    "metadata": {},
    "outputs": [],
    "source": [
diff --git a/examples/notebooks/Functional_API/constrained.ipynb b/examples/notebooks/Functional_API/constrained.ipynb
index 7350320..7e6f739 100644
--- a/examples/notebooks/Functional_API/constrained.ipynb
+++ b/examples/notebooks/Functional_API/constrained.ipynb
@@ -2,7 +2,7 @@
  "cells": [
   {
    "cell_type": "markdown",
-   "id": "deeeb625-f81b-44a3-99dc-732651fc1c01",
+   "id": "0",
    "metadata": {},
    "source": [
     "# Solving constrained optimization problems\n",
@@ -38,7 +38,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "82423f38-9832-4dab-9b0f-00893427576f",
+   "id": "1",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -59,7 +59,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "fd9964a8-0729-48f0-afec-9827a7cd1546",
+   "id": "2",
    "metadata": {},
    "source": [
     "## Fitness function implementation\n",
@@ -76,7 +76,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "4a1eddff-6e05-4037-b2eb-f37478734149",
+   "id": "3",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -131,7 +131,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "9a69f467-2838-41bb-bd93-f715255edf8f",
+   "id": "4",
    "metadata": {},
    "source": [
     "---\n",
@@ -141,7 +141,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "7d0d8808-0d9a-43f8-a275-df87fe3b6833",
+   "id": "5",
    "metadata": {},
    "source": [
     "As the evolutionary algorithm, we use the functional `pgpe`."
@@ -150,7 +150,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "b8ef5439-ee00-40e4-9053-cb621aa4c855",
+   "id": "6",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -194,7 +194,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "3474900a-1e39-4bbc-ad05-6ae377ac9143",
+   "id": "7",
    "metadata": {},
    "source": [
     "Main loop of the evolutionary computation:"
@@ -203,7 +203,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "d9c8ef56-117d-49f3-8dc2-9a77ae80af36",
+   "id": "8",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -243,7 +243,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "ec14d993-fd9a-4cd7-a772-741740aa0c7b",
+   "id": "9",
    "metadata": {},
    "source": [
     "Best solution:"
@@ -252,7 +252,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "b53479ed-2350-4fe3-ab4a-e2036aafce3d",
+   "id": "10",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -261,7 +261,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "20c93618-deb8-41dc-a06d-18dc78d96497",
+   "id": "11",
    "metadata": {},
    "source": [
     "---\n",
@@ -278,7 +278,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "3d12cfd7-25cb-4ff5-8481-dba1d8ec8b00",
+   "id": "12",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -321,7 +321,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "b09576f6-0a84-4572-a36f-9a6c4b3c7f4f",
+   "id": "13",
    "metadata": {},
    "source": [
     "---\n",
@@ -331,7 +331,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "b8d18eca-040e-4156-a26f-46af1375c886",
+   "id": "14",
    "metadata": {},
    "source": [
     "Although, as its name implies, the main focus of EvoTorch is evolutionary computation, we now demonstrate that it is possible to implement an interior points method by combining:\n",
@@ -348,7 +348,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "40a5bc1c-e6f2-4868-a426-5c24b4b62924",
+   "id": "15",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -384,7 +384,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "f4128f4f-d8ad-4be9-9d81-968406355f8f",
+   "id": "16",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -406,7 +406,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "19c0cfad-3ae3-4c3b-b6a9-ddc99051c3fe",
+   "id": "17",
    "metadata": {},
    "outputs": [],
    "source": [
diff --git a/examples/notebooks/Functional_API/problem.ipynb b/examples/notebooks/Functional_API/problem.ipynb
index e5103db..644a415 100644
--- a/examples/notebooks/Functional_API/problem.ipynb
+++ b/examples/notebooks/Functional_API/problem.ipynb
@@ -2,7 +2,7 @@
  "cells": [
   {
    "cell_type": "markdown",
-   "id": "e89116a2-a149-487c-8e8f-7dab78ded7ff",
+   "id": "0",
    "metadata": {},
    "source": [
     "# Solving reinforcement learning tasks using functional evolutionary algorithms\n",
@@ -13,7 +13,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "0286d0f6-8879-4722-902d-1caeb85306bc",
+   "id": "1",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -25,7 +25,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "eb10227b-523d-4ca5-a871-aca97063eba4",
+   "id": "2",
    "metadata": {},
    "source": [
     "Below, we instantiate the reinforcement learning problem."
@@ -34,7 +34,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "9af709e7-796a-47b1-8501-123e3f9fbec1",
+   "id": "3",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -65,7 +65,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "7d55d935-07fe-423a-b75b-befae8c97a21",
+   "id": "4",
    "metadata": {},
    "source": [
     "Now that we have instantiated our problem, we make a callable evaluator object from it. This callable evaluator, named `f`, behaves like a function `f(x)`, where `x` can be a single solution (represented by a 1-dimensional tensor), or a population (represented by a 2-dimensional tensor where each row is a solution), or a batch of populations (represented by a tensor with at least 3 dimensions). Upon receiving its argument `x`, `f` uses the problem object to evaluate the solution(s), and return the evalution result(s)."
@@ -74,7 +74,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "07619fad-0604-42ce-af08-3417917a3dd7",
+   "id": "5",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -84,7 +84,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "7ea2af7d-78ee-46e5-bfa4-f959b5565b3c",
+   "id": "6",
    "metadata": {},
    "source": [
     "Hyperparameters for `pgpe`:"
@@ -93,7 +93,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "9a95f546-1c15-4eed-9d2f-6fceb2d51a64",
+   "id": "7",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -108,7 +108,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "82c77b1e-318b-450d-9d3c-c5ad6cede329",
+   "id": "8",
    "metadata": {},
    "source": [
     "We prepare the `pgpe` algorithm and get its initial state:"
@@ -117,7 +117,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "df014e8e-a878-4ec1-80d4-392156e5ffbe",
+   "id": "9",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -161,7 +161,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "84e6a2be-c00d-4fee-b4f4-98b3156a90fb",
+   "id": "10",
    "metadata": {},
    "source": [
     "Below is the main loop of the evolutionary search."
@@ -170,7 +170,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "6adefd82-e819-49f4-b08d-7507d5efe264",
+   "id": "11",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -201,7 +201,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "7b2e117d-2bce-49ce-b912-745f5c707bf5",
+   "id": "12",
    "metadata": {},
    "source": [
     "Here is the center point of the most recent search distribution:"
@@ -210,7 +210,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "81a925c3-6cb2-4095-a01b-7c9ea275b575",
+   "id": "13",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -220,7 +220,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "c5e5215d-5ab4-41af-b421-db8684f9bfae",
+   "id": "14",
    "metadata": {},
    "source": [
     "Now, we visualize the agent evolved by our functional `pgpe`:"
@@ -229,7 +229,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "ce8321b8-7d87-430d-a144-2b0d3fd0d316",
+   "id": "15",
    "metadata": {},
    "outputs": [],
    "source": [
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",