diff --git a/burn-book/src/building-blocks/module.md b/burn-book/src/building-blocks/module.md
index 8397bc56ff..128a99a7b6 100644
--- a/burn-book/src/building-blocks/module.md
+++ b/burn-book/src/building-blocks/module.md
@@ -89,18 +89,29 @@ You can implement your own mapper or visitor by implementing these simple traits
 ```rust, ignore
 /// Module visitor trait.
 pub trait ModuleVisitor<B: Backend> {
-    /// Visit a tensor in the module.
-    fn visit<const D: usize>(&mut self, id: ParamId, tensor: &Tensor<B, D>);
+    /// Visit a float tensor in the module.
+    fn visit_float<const D: usize>(&mut self, id: ParamId, tensor: &Tensor<B, D>);
+    /// Visit an int tensor in the module.
+    fn visit_int<const D: usize>(&mut self, id: ParamId, tensor: &Tensor<B, D, Int>);
+    /// Visit a bool tensor in the module.
+    fn visit_bool<const D: usize>(&mut self, id: ParamId, tensor: &Tensor<B, D, Bool>);
 }
 
 /// Module mapper trait.
 pub trait ModuleMapper<B: Backend> {
-    /// Map a tensor in the module.
-    fn map<const D: usize>(&mut self, id: ParamId, tensor: Tensor<B, D>) ->
-      Tensor<B, D>;
+    /// Map a float tensor in the module.
+    fn map_float<const D: usize>(&mut self, id: ParamId, tensor: Tensor<B, D>) -> Tensor<B, D>;
+    /// Map an int tensor in the module.
+    fn map_int<const D: usize>(&mut self, id: ParamId, tensor: Tensor<B, D, Int>) -> Tensor<B, D, Int>;
+    /// Map a bool tensor in the module.
+    fn map_bool<const D: usize>(&mut self, id: ParamId, tensor: Tensor<B, D, Bool>) -> Tensor<B, D, Bool>;
 }
 ```
 
+Note that the trait doesn't require all methods to be implemented as they are already defined to
+perform no operation. If you're only interested in float tensors (like the majority of use cases),
+then you can simply implement `map_float` or `visit_float`.
+
 ## Module Display
 
 Burn provides a simple way to display the structure of a module and its configuration at a glance.
@@ -182,14 +193,14 @@ Burn comes with built-in modules that you can use to build your own modules.
 
 ### Convolutions
 
-| Burn API          | PyTorch Equivalent   |
-| ----------------- | -------------------- |
-| `Conv1d`          | `nn.Conv1d`          |
-| `Conv2d`          | `nn.Conv2d`          |
-| `Conv3d`          | `nn.Conv3d`          |
-| `ConvTranspose1d` | `nn.ConvTranspose1d` |
-| `ConvTranspose2d` | `nn.ConvTranspose2d` |
-| `ConvTranspose3d` | `nn.ConvTranspose3d` |
+| Burn API          | PyTorch Equivalent             |
+| ----------------- | ------------------------------ |
+| `Conv1d`          | `nn.Conv1d`                    |
+| `Conv2d`          | `nn.Conv2d`                    |
+| `Conv3d`          | `nn.Conv3d`                    |
+| `ConvTranspose1d` | `nn.ConvTranspose1d`           |
+| `ConvTranspose2d` | `nn.ConvTranspose2d`           |
+| `ConvTranspose3d` | `nn.ConvTranspose3d`           |
 | `DeformConv2d`    | `torchvision.ops.DeformConv2d` |
 
 ### Pooling