diff --git a/docs/source/layers.md b/docs/source/layers.md
index bf812f27..b2071cc6 100644
--- a/docs/source/layers.md
+++ b/docs/source/layers.md
@@ -65,10 +65,19 @@ compatible with layers from the hub.
### Using a function as a layer
-Sometimes it can be useful to make a function extensible, for example
-because the function cannot be replaced by a layer. In such cases, you
-can use the [`~kernels.use_kernel_forward_from_hub`] decorator on a
-function:
+
+> [!WARNING]
+> Decorating functions is mainly intended for compatibility cases, such as in
+`transformers`, where a function cannot be converted into a module. If you
+control the model code, prefer converting the function to an `nn.Module` and
+decorating that class with [`~kernels.use_kernel_forward_from_hub`], as shown
+in [Using a decorator](#using-a-decorator). In the parent module's `__init__`,
+assign an instance of the decorated class to an attribute, such as
+`self.activation = SiluAndMul()`. PyTorch then registers it as a submodule,
+making it part of the model graph.
+
+When a function cannot be replaced by a layer, you can use the
+[`~kernels.use_kernel_forward_from_hub`] decorator on the function:
```python
@use_kernel_forward_from_hub("silu_and_mul")
diff --git a/kernels/src/kernels/layer/layer.py b/kernels/src/kernels/layer/layer.py
index 80dd58ae..c22050e4 100644
--- a/kernels/src/kernels/layer/layer.py
+++ b/kernels/src/kernels/layer/layer.py
@@ -373,6 +373,18 @@ def use_kernel_forward_from_hub(layer_name: str, condition: Callable[["nn.Module
used as if it was a function. Note that a decorated function is only visible to
[`~kernels.kernelize`] if it is attached to a module using [`~kernels.use_kernelized_func`].
+
+
+ Decorating functions is mainly intended for compatibility cases, such as in
+ `transformers`, where a function cannot be converted into a module. If you
+ control the model code, prefer converting the function to an `nn.Module` and
+ decorating that class with [`~kernels.use_kernel_forward_from_hub`]. In the
+ parent module's `__init__`, assign an instance of the decorated class to an
+ attribute, such as `self.activation = MyActivation()`. PyTorch then registers
+ it as a submodule, making it part of the model graph.
+
+
+
Args:
layer_name (`str`):
The name of the layer to use for kernel lookup in registered mappings.
@@ -484,6 +496,18 @@ def use_kernelized_func(*args: Callable):
attribute (not as a submodule). This makes the function visible to
[`~kernels.kernelize`].
+
+
+ This decorator is mainly intended for compatibility cases, such as in
+ `transformers`, where a function cannot be converted into a module. If you
+ control the model code, prefer converting the function to an `nn.Module` and
+ decorating that class with [`~kernels.use_kernel_forward_from_hub`]. In the
+ parent module's `__init__`, assign an instance of the decorated class to an
+ attribute, such as `self.activation = MyActivation()`. PyTorch then registers
+ it as a submodule, making it part of the model graph.
+
+
+
Args:
*args (`Callable`):
Kernel functions to attach to the module.