diff --git a/predict_zaugnet.py b/predict_zaugnet.py
index 8160b4d..b614a08 100644
--- a/predict_zaugnet.py
+++ b/predict_zaugnet.py
@@ -21,7 +21,7 @@ def load_model(cfg, dataset, model_name):
 
     zaug = ZAugGenerator(cfg)
     zaug.set_multiple_gpus()
-    zaug.load_state_dict(torch.load(path_model))
+    zaug.load_state_dict(torch.load(path_model,  map_location=f"cuda:{cfg.device_ids[0]}"))
     zaug.eval()
 
     return zaug
diff --git a/predict_zaugnet_plus.py b/predict_zaugnet_plus.py
index d2a9c6c..a312a87 100644
--- a/predict_zaugnet_plus.py
+++ b/predict_zaugnet_plus.py
@@ -19,7 +19,7 @@ def load_model(cfg, dataset, model_name):
     print(f"The model used for this prediction : {path_model}")
     zaug = ZAugGenerator(cfg)
     zaug.set_multiple_gpus()
-    zaug.load_state_dict(torch.load(path_model))
+    zaug.load_state_dict(torch.load(path_model,  map_location=f"cuda:{cfg.device_ids[0]}"))
     zaug.eval()
 
     return zaug
diff --git a/zaugnet_colab.ipynb b/zaugnet_colab.ipynb
index 0d1d08d..6690f23 100644
--- a/zaugnet_colab.ipynb
+++ b/zaugnet_colab.ipynb
@@ -88,7 +88,6 @@
       "outputs": [],
       "source": [
         "!git clone https://github.com/VirtualEmbryo/ZAugNet.git"
-<<<<<<< HEAD
       ]
     },
     {
@@ -122,43 +121,8 @@
         "!mv /content/ZAugNet/zaugnet_data_and_models/zenodo/humans /content/ZAugNet/data/humans\n",
         "!mv /content/ZAugNet/zaugnet_data_and_models/zenodo/nuclei /content/ZAugNet/data/nuclei\n",
         "!mv /content/ZAugNet/zaugnet_data_and_models/zenodo/results /content/ZAugNet/results"
-=======
->>>>>>> 24dd0f40c0f345e5b34afd3b910460dc071e1e15
       ]
     },
-    {
-      "cell_type": "markdown",
-      "source": [
-        "* **(Optional)** Download train/test data images and pre-trained models:"
-      ],
-      "metadata": {
-        "id": "WdKYdfzX2nge"
-      }
-    },
-    {
-      "cell_type": "code",
-      "source": [
-        "# Download data and pre-trained models from Zenodo\n",
-        "!curl \"https://zenodo.org/records/14961732/files/zenodo.zip?download=1\" --output /content/ZAugNet/zaugnet_data_and_models.zip\n",
-        "\n",
-        "# Unzip the folder\n",
-        "import shutil\n",
-        "shutil.unpack_archive(\"/content/ZAugNet/zaugnet_data_and_models.zip\", \"/content/ZAugNet/zaugnet_data_and_models\")\n",
-        "\n",
-        "# Create data folder and move files\n",
-        "!mkdir /content/ZAugNet/data\n",
-        "!mv /content/ZAugNet/zaugnet_data_and_models/zenodo/ascidians /content/ZAugNet/data/ascidians\n",
-        "!mv /content/ZAugNet/zaugnet_data_and_models/zenodo/filaments /content/ZAugNet/data/filaments\n",
-        "!mv /content/ZAugNet/zaugnet_data_and_models/zenodo/humans /content/ZAugNet/data/humans\n",
-        "!mv /content/ZAugNet/zaugnet_data_and_models/zenodo/nuclei /content/ZAugNet/data/nuclei\n",
-        "!mv /content/ZAugNet/zaugnet_data_and_models/zenodo/results /content/ZAugNet/results"
-      ],
-      "metadata": {
-        "id": "osmjfhvc23lq"
-      },
-      "execution_count": null,
-      "outputs": []
-    },
     {
       "cell_type": "markdown",
       "metadata": {
@@ -456,4 +420,4 @@
   },
   "nbformat": 4,
   "nbformat_minor": 0
-}
\ No newline at end of file
+}