diff --git a/train.py b/train.py
index 170dac6..b8b1de3 100644
--- a/train.py
+++ b/train.py
@@ -28,7 +28,7 @@
 from urllib.request import urlretrieve
 import shutil
 
-DATASET_VERSION = 'v2.0'
+DATASET_VERSION = 'v2.1'
 DATASET_ZIP_URL = f'https://github.com/pynicolas/fairscan-dataset/releases/download/{DATASET_VERSION}/fairscan-dataset-{DATASET_VERSION}.zip'
 
 BUILD_DIR = "build"
@@ -38,7 +38,7 @@
 DATASET_ZIP_PATH = BUILD_DIR + "/dataset.zip"
 DATASET_PARENT_DIR = BUILD_DIR + "/dataset"
 DATASET_DIR = DATASET_PARENT_DIR + "/fairscan-dataset"
-NB_EPOCHS = 25
+NB_EPOCHS = 35
 
 if os.path.isdir(BUILD_DIR):
     shutil.rmtree(BUILD_DIR)