Best practice for after_action?

In after_action change, I usually load some relationships and call some update actions.
And the action would return a dirty record with old relationships.

Example

defmodule A do
  actions do
    update :update do
      change after_action(fn changeset, record, context ->
        record = Ash.load(record, [:b])
        B.update!(record.b, params)
        {:ok, record}
      end)
    end
  end
end

A.update action calls B.update in after_action.
The result of B.update is not saved anywhere but db.
Therefore, A.update action would return updated A record and old B record.
Should I always call Ash.reload!(record) at the end of after_action?
Or, is there any recommendation how to use after_action?

This is problematic when calling Ash.can with returned A record.