Use old custom grad

Author: seanmor5

lib/axon/compiler.ex

@@ -889,7 +889,7 @@ defmodule Axon.Compiler do
 
       if event? and mode? do
         if on_event == :backward do
-          Nx.Defn.Kernel.custom_grad(expr, [expr], fn g ->
+          Nx.Defn.Kernel.custom_grad(expr, fn _ans, g ->
             hooked_g = Nx.Defn.Kernel.hook(g, hook_fn)
             [{expr, hooked_g}]
           end)