make attach_hook API compatible with Axon.Node and fix doc example on…

… Axon.map_nodes (#519)

lib/axon.ex

@@ -3177,19 +3177,25 @@ defmodule Axon do
 
   """
   @doc type: :debug
-  def attach_hook(%Axon{output: id, nodes: nodes} = axon, fun, opts \\ []) do
-    opts = Keyword.validate!(opts, on: :forward, mode: :both)
-    on_event = opts[:on]
-    mode = opts[:mode]
+  def attach_hook(x, fun, opts \\ [])
 
+  def attach_hook(%Axon{output: id, nodes: nodes} = axon, fun, opts) do
     updated_nodes =
       Map.update!(nodes, id, fn axon_node ->
-        %{axon_node | hooks: [{on_event, mode, fun} | axon_node.hooks]}
+        attach_hook(axon_node, fun, opts)
       end)
 
     %{axon | nodes: updated_nodes}
   end
 
+  def attach_hook(%Axon.Node{hooks: hooks} = axon_node, fun, opts) do
+    opts = Keyword.validate!(opts, on: :forward, mode: :both)
+    on_event = opts[:on]
+    mode = opts[:mode]
+
+    %{axon_node | hooks: [{on_event, mode, fun} | hooks]}
+  end
+
   ## Graph Manipulation and Utilities
 
   # TODO: Revisit later with new decoupled structs
@@ -3353,12 +3359,12 @@ defmodule Axon do
   you can use this function to visualize intermediate activations
   of all convolutional layers in a model:
 
-      instrumented_model = Axon.  (model, fn
-        %Axon{op: :conv} = graph ->
-          Axon.attach_hook(graph, &visualize_activations/1)
+      instrumented_model = Axon.map_nodes(model, fn
+        %Axon.Node{op: :conv} = axon_node ->
+          Axon.attach_hook(axon_node, &visualize_activations/1)
 
-        graph ->
-          graph
+        axon_node ->
+          axon_node
       end)
 
   Another use case is to replace entire classes of layers