Efficiently resolving belongs_to relationship in elixir dataloader?

Is it possible to use elixir dataloader to query a belongs_to relationship efficiently? It seems that the load is querying all of the items it needs, but the get is returning the first value of the loaded items regardless of which single item it actually needs. This is the code that I am using now:

field :node, :node_object, resolve: fn parent, _, %{context: %{loader: loader}} ->
  # parent.node_id = 1, but concurrently also another parent.node_id = 5
  loader
  |> Dataloader.load(NodeContext, :node, parent) # loads node_id 5 and 1
  |> on_load(fn loader ->
    loader
    |> Dataloader.get(NodeContext, :node, parent) # always returns the node with id = 5
    |> (&{:ok, &1}).()
  end)                                                     
end

My current work around is to use the following code, but it makes the code much uglier and unfriendly with the Ecto schemas since I need to explicitly specify the node schema and node_id field of the parent schema here instead of letting dataloader infer it from the existing ecto schemas:

field :node, :node_object, resolve: fn parent, _, %{context: %{loader: loader}} ->
  loader
  |> Dataloader.load(NodeContext, {:one, NodeSchema}, id: parent.node_id)
  |> on_load(fn loader ->
    loader
    |> Dataloader.get(NodeContext, {:one, NodeSchema}, id: parent.node_id)
    |> (&{:ok, &1}).()
  end)
end

I was able to fix this by making the node_id a primary_key of the parent schema like this:

defmodule MyApp.ParentSchema do
  use Ecto.Schema
  alias MyApp.NodeSchema

  @primary_key false

  embedded_schema do
    belongs_to :node, NodeSchema, primary_key: true
  end
end

I'm not sure if this is intended behavior for the dataloader since it seems like the primary_key check should happen on the child object instead of the parent object.