Added clarification to apply_fn parameter of apply_to_modules (#2831)

Co-authored-by: Awni Hannun <[email protected]>

python/mlx/nn/layers/base.py

@@ -407,7 +407,10 @@ class Module(dict):
         instance).
 
         Args:
-            apply_fn (Callable): The function to apply to the modules.
+            apply_fn (Callable): The function to apply to the modules which
+                takes two parameters. The first parameter is the string path of
+                the module (e.g. ``"model.layers.0.linear"``). The second
+                parameter is the module object.
 
         Returns:
             The module instance after updating submodules.