diff --git a/colab/Colab_UnityEnvironment_1_Run.ipynb b/colab/Colab_UnityEnvironment_1_Run.ipynb
index e1af12ef5e..c423ddd6f1 100644
--- a/colab/Colab_UnityEnvironment_1_Run.ipynb
+++ b/colab/Colab_UnityEnvironment_1_Run.ipynb
@@ -536,7 +536,7 @@
     "for index, obs_spec in enumerate(spec.observation_specs):\n",
     "  if len(obs_spec.shape) == 3:\n",
     "    print(\"Here is the first visual observation\")\n",
-    "    plt.imshow(np.rollaxis(decision_steps.obs[index][0, :, :, :], 0, 3))\n",
+    "    plt.imshow(np.moveaxis(decision_steps.obs[index][0, :, :, :], 0, -1))\n",
     "    plt.show()\n",
     "\n",
     "for index, obs_spec in enumerate(spec.observation_specs):\n",
diff --git a/colab/Colab_UnityEnvironment_3_SideChannel.ipynb b/colab/Colab_UnityEnvironment_3_SideChannel.ipynb
index 3eb368bae5..9d7cfe7d4d 100644
--- a/colab/Colab_UnityEnvironment_3_SideChannel.ipynb
+++ b/colab/Colab_UnityEnvironment_3_SideChannel.ipynb
@@ -326,7 +326,7 @@
     "\n",
     "print(\"Observation without changing the environment parameters\")\n",
     "decision_steps, terminal_steps = env.get_steps(behavior_name)\n",
-    "plt.imshow(np.rollaxis(decision_steps.obs[0][0,:,:,:], 0, 3))\n",
+    "plt.imshow(np.moveaxis(decision_steps.obs[0][0,:,:,:], 0, -1))\n",
     "plt.show()\n",
     "\n",
     "print(\"Increasing the dimensions of the grid from 5 to 7\")\n",
@@ -340,7 +340,7 @@
     "\n",
     "\n",
     "decision_steps, terminal_steps = env.get_steps(behavior_name)\n",
-    "plt.imshow(np.rollaxis(decision_steps.obs[0][0,:,:,:], 0, 3))\n",
+    "plt.imshow(np.moveaxis(decision_steps.obs[0][0,:,:,:], 0, -1))\n",
     "plt.show()\n",
     "\n",
     "\n",