Add a basic inference mode (elixir-nx#129)

zblanco · Oct 7, 2021 · 4273e64 · 4273e64
1 parent d42cc99
commit 4273e64
Show file tree

Hide file tree

Showing 4 changed files with 115 additions and 62 deletions.
diff --git a/lib/axon.ex b/lib/axon.ex
@@ -51,7 +51,11 @@ defmodule Axon do
 
       params = Axon.init(model, compiler: EXLA)
 
-      Axon.predict(model, params, inputs, compiler: EXLA)
+      Axon.predict(model, params, inputs, compiler: EXLA, mode: :train)
+
+  `Axon.predict/4` by default runs in inference mode, which performs certain
+  optimizations and removes layers such as dropout layers. If constructing
+  a training step using `Axon.predict/4`, be sure to specify `mode: :train`.
 
   Both `Axon.init/2` and `Axon.predict/4` can be used from within
   Nx defn or outside.