make the example run (#37)

Author: mcabbott

docs/src/index.md

@@ -40,17 +40,19 @@ to adjust the model:
 
 using Flux, Metalhead, Optimisers
 
-o = Optimisers.ADAM() # define an ADAM optimiser with default settings
-st = Optimisers.setup(o, m)  # initialize the optimiser before using it
+model = Metalhead.ResNet18() # define a model to train on
+image = rand(Float32, 224, 224, 3, 1); # dummy data
+@show sum(model(image)); # dummy loss function
 
-model = ResNet18() # define a model to train on
-ip = rand(Float32, 224, 224, 3, 1) # dummy data
+o = Optimisers.ADAM() # define an ADAM optimiser with default settings
+st = Optimisers.setup(o, model);  # initialize the optimiser before using it
 
-m̄, _ = gradient(model, ip) do m, x # calculate the gradients
-  sum(m(x)) # dummy loss function
-end
+m̄, _ = gradient(model, image) do m, x # calculate the gradients
+  sum(m(x))
+end;
 
-st, mnew = Optimisers.update(st, m, m̄)
+st, model = Optimisers.update(st, model, m̄);
+@show sum(model(image));
 
 ```