Error when using before_action hook to set required relationship attribute

create :create do
  primary? true
  upsert? true
  upsert_identity :unique_google_place
  upsert_fields {:replace_all_except, [:id, :state_id, :inserted_at]}
  skip_unknown_inputs :*

  change __MODULE__.Changes.SetState,
    where: [present(:google_state), absent(:state_id)]
end

defmodule Boring.Google.Place.Changes.SetState do
  @moduledoc false
  use Ash.Resource.Change

  require Logger

  @impl Ash.Resource.Change
  def change(changeset, _opts, context) do
    Ash.Changeset.before_action(changeset, fn changeset ->
      state = Ash.Changeset.get_argument_or_attribute(changeset, :google_state)
      Logger.debug("Setting state for #{state}")

      opts = Ash.Context.to_opts(context)

      case Boring.Locations.get_state_by_name_or_code(state, opts) do
        %Boring.Locations.State{} = matching_state ->
          Ash.Changeset.force_change_attribute(changeset, :state_id, matching_state.id)

        _ ->
          changeset
      end
    end)
  end
end