diff --git a/stemflow/utils/lazyloading.py b/stemflow/utils/lazyloading.py
index 7ca66e6..c668ceb 100644
--- a/stemflow/utils/lazyloading.py
+++ b/stemflow/utils/lazyloading.py
@@ -7,7 +7,10 @@ class LazyLoadingEnsembleDict(MutableMapping):
     def __init__(self, directory='./tmp_models'):
         """
         Initialize the LazyLoadingEnsembleDict with a directory to save and load models.
-        :param directory: The directory to save and load model files.
+        
+        Args:
+            directory:
+                The directory to save and load model files.
         """
         self.directory = directory
         self.ensemble_models = {}
@@ -197,4 +200,4 @@ def delete_ensemble(self, ensemble_id):
         if os.path.exists(ensemble_path):
             os.remove(ensemble_path)
         else:
-            raise ValueError(f'Ensemble {ensemble_id} not found on disk.')
\ No newline at end of file
+            raise ValueError(f'Ensemble {ensemble_id} not found on disk.')