fix: removed the jax import and defaulted jax to use the cpu due to o…

…ut of memory issues in the unet demo (#72)

examples_and_demos/image_segmentation_with_ivy_unet.ipynb

@@ -61,8 +61,6 @@
       "source": [
         "import ivy\n",
         "ivy.set_default_device(\"gpu:0\")\n",
-        "import jax\n",
-        "jax.devices()\n",
         "import torch\n",
         "import numpy as np"
       ]
@@ -562,6 +560,7 @@
         "import jax\n",
         "\n",
         "jax.config.update('jax_enable_x64', True)\n",
+        "ivy.set_default_device(\"cpu\")\n",
         "ivy.set_backend(\"jax\")\n",
         "ivy_unet = ivy_models.unet_carvana(n_channels=3, n_classes=2, pretrained=True)"
       ]