revise readme

Author: t-rutten

README.md

@@ -2,12 +2,20 @@
 
 ## Usage
 
+We currently support MNIST, FashionMNIST, and CIFAR10 training datasets.
+
+Download or fetch datasets locally like:
+
 ```elixir
 {train_images, train_labels} = AxonDatasets.MNIST.download()
 
+# Unpack train_images like...
 {images_binary, tensor_type, shape} = train_images
+```
 
-# Transform the data as it's fetched
+You can also pass transform functions to `download/1`:
+
+```elixir
 transform_images = fn {binary, type, shape} ->
   binary
   |> Nx.from_binary(type)
@@ -20,7 +28,7 @@ end
   AxonDatasets.MNIST.download(transform_images: transform_images)
 
 # Transform labels as well, e.g. get one-hot encoding
-transform_labels = fn {labels_binary, type, _shape} ->
+transform_labels = fn {labels_binary, type, _} ->
   labels_binary
   |> Nx.from_binary(type)
   |> Nx.new_axis(-1)
@@ -45,7 +53,3 @@ def deps do
   ]
 end
 ```
-
-Documentation can be generated with [ExDoc](https://github.com/elixir-lang/ex_doc)
-and published on [HexDocs](https://hexdocs.pm). Once published, the docs can
-be found at [https://hexdocs.pm/axon_data](https://hexdocs.pm/axon_data).