It's not possible so far via Bumblebee, unfortunately[1].

[1] https://github.com/elixir-nx/bumblebee/issues/209