diff --git a/luxonis_ml/data/__init__.py b/luxonis_ml/data/__init__.py
index ebb85c11..66ba4505 100644
--- a/luxonis_ml/data/__init__.py
+++ b/luxonis_ml/data/__init__.py
@@ -22,14 +22,14 @@ def load_dataset_plugins() -> None:  # pragma: no cover
     """Registers any external dataset BaseDataset class plugins."""
     for entry_point in pkg_resources.iter_entry_points("dataset_plugins"):
         plugin_class = entry_point.load()
-        DATASETS_REGISTRY.register_module(module=plugin_class)
+        DATASETS_REGISTRY.register(module=plugin_class)
 
 
 def load_loader_plugins() -> None:  # pragma: no cover
     """Registers any external dataset BaseLoader class plugins."""
     for entry_point in pkg_resources.iter_entry_points("loader_plugins"):
         plugin_class = entry_point.load()
-        DATASETS_REGISTRY.register_module(module=plugin_class)
+        DATASETS_REGISTRY.register(module=plugin_class)
 
 
 load_dataset_plugins()
diff --git a/luxonis_ml/data/augmentations/custom/__init__.py b/luxonis_ml/data/augmentations/custom/__init__.py
index 55f6912c..1c565630 100644
--- a/luxonis_ml/data/augmentations/custom/__init__.py
+++ b/luxonis_ml/data/augmentations/custom/__init__.py
@@ -12,8 +12,8 @@
     "albumentations_transformations"
 )
 
-TRANSFORMATIONS.register_module(module=LetterboxResize)
-TRANSFORMATIONS.register_module(module=MixUp)
-TRANSFORMATIONS.register_module(module=Mosaic4)
+TRANSFORMATIONS.register(module=LetterboxResize)
+TRANSFORMATIONS.register(module=MixUp)
+TRANSFORMATIONS.register(module=Mosaic4)
 
 __all__ = ["LetterboxResize", "MixUp", "Mosaic4", "TRANSFORMATIONS"]
diff --git a/luxonis_ml/utils/registry.py b/luxonis_ml/utils/registry.py
index 49d4f9bd..ef68e608 100644
--- a/luxonis_ml/utils/registry.py
+++ b/luxonis_ml/utils/registry.py
@@ -76,7 +76,8 @@ def register_module(
         force: bool = False,
     ) -> Union[T, Callable[[T], T]]:
         warnings.warn(
-            "`register_module` is deprecated, use `register` instead."
+            "`register_module` is deprecated, use `register` instead.",
+            stacklevel=2,
         )
 
         return self.register(name=name, module=module, force=force)
@@ -111,14 +112,14 @@ def register(
         Can be used as a decorator or as a normal method:
 
             >>> registry = Registry(name="modules")
-            >>> @registry.register_module()
+            >>> @registry.register()
             ... class Foo:
             ...     pass
             >>> registry.get("Foo")
             <class '__main__.Foo'>
             >>> class Bar:
             ...     pass
-            >>> registry.register_module(module=Bar)
+            >>> registry.register(module=Bar)
             >>> registry.get("Bar")
             <class '__main__.Bar'>
 
@@ -227,7 +228,5 @@ def __new__(
                 )
         if register:
             registry = registry if registry is not None else new_class.REGISTRY
-            registry.register_module(
-                name=register_name or name, module=new_class
-            )
+            registry.register(name=register_name or name, module=new_class)
         return new_class