docs: document new loader_fun option

elixir-nx · Aug 3, 2023 · fbbf02b · fbbf02b
1 parent ef42ae6
commit fbbf02b
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/lib/bumblebee/conversion/pytorch.ex b/lib/bumblebee/conversion/pytorch.ex
@@ -24,6 +24,10 @@ defmodule Bumblebee.Conversion.PyTorch do
       between the Axon model and the PyTorch state. For more details see
       `Bumblebee.HuggingFace.Transformers.Model.params_mapping/1`
 
+    * `:loader_fun` - a 1-arity function that takes a path argument
+      and loads the params file.
+      Defaults to `Bumblebee.Conversion.PyTorch.Loader.load!/1`.
+
   """
   @spec load_params!(Axon.t(), map(), Path.t() | list(Path.t()), keyword()) :: map()
   def load_params!(model, input_template, path, opts \\ []) do