diff --git a/README.md b/README.md
index ba5ffc0..b8fbb0d 100644
--- a/README.md
+++ b/README.md
@@ -53,7 +53,12 @@ python code/predictor.py --prediction-path "prediction.tif" \\
 For the 5-fold training run
 ```
 python code/main.py train --results-dir models --data-path /data
-``` 
+```
+
+```
+!python code/train.py --data-path /content/floatingobjects/data --snapshot-path /content/floatingobjects/models/model.pth.tar --image-size 128 --epochs 50 --tensorboard-logdir /content/floatingobjects/models/runs/
+
+```
 
 #### Load pretrained models using PyTorch Hub
 Here you can load the models with pretrained weights from the 2-fold training of U-Net and MA-Net, according to what has been published in [OCEANS Conference Paper](https://210507-004.oceansvirtual.com/view/content/skdwP611e3583eba2b/ecf65c2aaf278557ad05c213247d67a54196c9376a0aed8f1875681f182daeed)
diff --git a/code/data.py b/code/data.py
index 0ae615f..177e73a 100644
--- a/code/data.py
+++ b/code/data.py
@@ -18,26 +18,17 @@
     "accra_20181031",
     "biscay_20180419",
     "danang_20181005",
-    "kentpointfarm_20180710",
     "kolkata_20201115",
     "lagos_20190101",
     "lagos_20200505",
-    "london_20180611",
     "longxuyen_20181102",
     "mandaluyong_20180314",  
     "neworleans_20200202",
     "panama_20190425",
-    "portalfredSouthAfrica_20180601",
     "riodejaneiro_20180504",
     "sandiego_20180804",
-    "sanfrancisco_20190219", 
     "shengsi_20190615",
-    "suez_20200403",
-    "tangshan_20180130",
     "toledo_20191221",
-    "tungchungChina_20190922",
-    "tunisia_20180715",
-    "turkmenistan_20181030",
     "venice_20180630",
     "venice_20180928",
     "vungtau_20180423"
@@ -249,4 +240,4 @@ def line_is_closed(linestringgeometry):
     coordinates = np.stack(linestringgeometry.xy).T
     first_point = coordinates[0]
     last_point = coordinates[-1]
-    return bool((first_point == last_point).all())
+    return bool((first_point == last_point).all())
\ No newline at end of file
diff --git a/code/download.py b/code/download.py
index 760187b..c38c55e 100644
--- a/code/download.py
+++ b/code/download.py
@@ -12,13 +12,20 @@
     parser.add_argument('storepath', type=str, nargs='?', default=os.path.join(this_folder, "..", "data"))
     args = parser.parse_args()
 
-    print(f"creating {os.path.abspath(args.storepath)}")
+    # Ensure the directory exists
+    print(f"Creating {os.path.abspath(args.storepath)}")
     os.makedirs(args.storepath, exist_ok=True)
 
-    # download
-    print(f"downloading {URL} to {os.path.abspath(args.storepath)}")
-    gdown.download(URL, args.storepath, quiet=False)
+    # Specify the path including the filename where the file will be saved
+    zip_path = os.path.join(args.storepath, "data.zip")
 
-    # unpack
-    print(f"unpacking {args.storepath}/data.zip to {os.path.abspath(args.storepath)}")
-    shutil.unpack_archive(os.path.join(args.storepath, "data.zip"), os.path.join(args.storepath, ".."))
\ No newline at end of file
+    # Download
+    print(f"Downloading {URL} to {zip_path}")
+    gdown.download(URL, zip_path, quiet=False)
+
+    # Unpack
+    print(f"Unpacking {zip_path} to {os.path.abspath(args.storepath)}")
+    shutil.unpack_archive(zip_path, os.path.abspath(args.storepath))
+
+    # Optionally, you can delete the zip file after unpacking if it's no longer needed
+    # os.remove(zip_path)