Fix missing state in training (#579)

* Inspect

* Inspect more

* Raise on key

* inspect

* Fix?

* Again

* Again

* Fix axon

* It works

lib/axon/loop.ex

@@ -1556,7 +1556,7 @@ defmodule Axon.Loop do
       is set, the loop will raise on any cache miss during the training loop. Defaults
       to true.
 
-    * `:force_garbage_collect?` - whether or not to force garbage collection after each
+    * `:force_garbage_collection?` - whether or not to force garbage collection after each
       iteration. This may help avoid OOMs when training large models, but it will slow
       training down.
 

lib/axon/model_state.ex

@@ -222,13 +222,27 @@ defmodule Axon.ModelState do
   end
 
   defp tree_get(data, access) when is_list(access) do
-    Enum.reduce(access, %{}, &Map.put(&2, &1, Map.fetch!(data, &1)))
+    Enum.reduce(access, %{}, fn key, acc ->
+      case data do
+        %{^key => val} ->
+          Map.put(acc, key, val)
+
+        %{} ->
+          acc
+      end
+    end)
   end
 
   defp tree_get(data, access) when is_map(access) do
     Enum.reduce(access, %{}, fn {key, value}, acc ->
-      tree = tree_get(data[key], value)
-      Map.put(acc, key, tree)
+      case data do
+        %{^key => val} ->
+          tree = tree_get(val, value)
+          Map.put(acc, key, tree)
+
+        %{} ->
+          acc
+      end
     end)
   end