Fix prepare_dataset issue

Author: jrrodri

abraia/__init__.py

@@ -2,7 +2,7 @@
 from dotenv import load_dotenv
 load_dotenv()
 
-__version__ = '0.25.5'
+__version__ = '0.25.6'
 
 from . import config
 from .client import Abraia, APIError

abraia/training/__init__.py

@@ -6,8 +6,8 @@
 from typing import Dict, Any
 from tqdm.contrib.concurrent import process_map
 
-from .dataset import list_datasets, load_dataset, search_images, list_models
-from ..utils import make_dirs, download_file
+from .dataset import list_datasets, load_dataset, search_images, list_models, download_file
+from ..utils import make_dirs
 
 
 def load_tasks(task):

abraia/training/dataset.py

@@ -247,7 +247,7 @@ def load(self):
             self.classes = load_labels(self.annotations)
             self.task = load_task(self.annotations)
             self.images = list_images(self.project)
-        return self
+            return self
 
     def annotate(self, label, segment=False):
         annotated_filenames = {a['filename'] for a in self.annotations}