Using bumblebee from git is broken

[dennym]

Since there where minor improvements on github since v0.5.3 and no new released version I'm using bumblebee from git in my mix.exs like {:bumblebee, git: "https://github.com/elixir-nx/bumblebee"}. This seems currently broken and when trying to .load_model/2 it fails due to some Axon changes which I can not debug where the change might come from:

iex(3)> {:ok, model_info} = Bumblebee.load_model(repo, backend: {EXLA.Backend, client: :host})
[warning] passing parameter map to initialization is deprecated, use %Axon.ModelState{} instead
** (UndefinedFunctionError) function Axon.ModelState.fetch/2 is undefined (Axon.ModelState does not implement the Access behaviour

You can use the "struct.field" syntax to access struct fields. You can also use Access.key!/1 to access struct fields dynamically inside get_in/put_in/update_in)
    (axon 0.6.1) Axon.ModelState.fetch(%Axon.ModelState{data: %{....

Erlang/OTP 26 [erts-14.2.4] [source] [64-bit] [smp:6:6] [ds:6:6:10] [async-threads:1] [jit:ns]

Interactive Elixir (1.16.2) - press Ctrl+C to exit (type h() ENTER for help)

Any ideas?

[seanmor5]

For now to fix you can pin Axon to 0.6.1 while I work on fixing this

[jonatanklosko]

Try {:axon, github: "elixir-nx/axon", ref: "7e0e5930ac4b8d2a89f48106b8121e103e597c89", force: true} (that's the revision we have in mix.lock and it should work).