Prep for release 0.6 (#521)

Author: seanmor5

README.md

@@ -125,7 +125,7 @@ Then add Axon to your dependencies:
 ```elixir
 def deps do
   [
-    {:axon, "~> 0.5"}
+    {:axon, "~> 0.6"}
   ]
 end
 ```
@@ -135,8 +135,8 @@ You'll also likely want to include an `Nx` compiler such as `EXLA` for any pract
 ```elixir
 def deps do
   [
-    {:axon, "~> 0.5"},
-    {:exla, "~> 0.5"},
+    {:axon, "~> 0.6"},
+    {:exla, "~> 0.6"},
   ]
 end
 ```

mix.exs

@@ -2,7 +2,7 @@ defmodule Axon.MixProject do
   use Mix.Project
 
   @source_url "https://github.com/elixir-nx/axon"
-  @version "0.5.1"
+  @version "0.6.0"
 
   def project do
     [
@@ -35,9 +35,9 @@ defmodule Axon.MixProject do
   # Run "mix help deps" to learn about dependencies.
   defp deps do
     [
-      {:exla, "~> 0.5.0", [only: :test] ++ exla_opts()},
-      {:torchx, "~> 0.5.0", [only: :test] ++ torchx_opts()},
-      {:nx, "~> 0.5.0", nx_opts()},
+      {:exla, "~> 0.6.0", [only: :test] ++ exla_opts()},
+      {:torchx, "~> 0.6.0", [only: :test] ++ torchx_opts()},
+      {:nx, "~> 0.6.0", nx_opts()},
       {:ex_doc, "~> 0.23", only: :docs},
       {:table_rex, "~> 3.1.1", optional: true},
       {:kino, "~> 0.7", optional: true},