diff --git a/training/training/core/dataset.py b/training/training/core/dataset.py
index 4efe314a0..6a28b9c8f 100644
--- a/training/training/core/dataset.py
+++ b/training/training/core/dataset.py
@@ -161,7 +161,7 @@ def fromDefault(
         test_transform=None,
         batch_size: int = 32,
         shuffle: bool = True,
-    ):
+    ) -> "ImageDefaultDatasetCreator":
         return cls(dataset_name, train_transform, test_transform, batch_size, shuffle)
     
     def delete_datasets_from_directory(self):
diff --git a/training/training/routes/image/__init__.py b/training/training/routes/image/__init__.py
index d24687888..89149b64b 100644
--- a/training/training/routes/image/__init__.py
+++ b/training/training/routes/image/__init__.py
@@ -1,5 +1,5 @@
 from training.routes.image.image import router
 
 
-def get_image_router():
+def get_image_router() -> router:
     return router