Skip to content

Fix DinoV2 crash when batch_size > 1 - #429

Merged
jonatanklosko merged 3 commits into
elixir-nx:mainfrom
zacharydenton:main
Dec 18, 2025
Merged

Fix DinoV2 crash when batch_size > 1#429
jonatanklosko merged 3 commits into
elixir-nx:mainfrom
zacharydenton:main

Conversation

@zacharydenton

Copy link
Copy Markdown
Contributor

This fixes an error in interpolate_position_embeddings when batch_size is > 1 (#428).

Comment thread lib/bumblebee/vision/dino_v2.ex Outdated
interpolated_position_embeddings =
input_position_embeddings
|> Nx.reshape({batch_size, original_positions, original_positions, spec.hidden_size})
|> Nx.reshape({1, original_positions, original_positions, spec.hidden_size})

@jonatanklosko jonatanklosko Dec 18, 2025

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@zacharydenton great find! I would suggest a small improvement:

position_embeddings_batch_size = Nx.axis_size(position_embeddings, 0)

and then use that in both places.

In practice it's always 1, but hardcoding the value implies an extra assumption that the reader may need to understand, which means going back to the caller side and figuring out the position embeddings shape :)

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done.

Comment thread lib/bumblebee/vision/dino_v2.ex Outdated

@jonatanklosko jonatanklosko left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thank you!

@jonatanklosko jonatanklosko changed the title Fix DinoV2 crash when batch_size > 1. Fix DinoV2 crash when batch_size > 1 Dec 18, 2025
@jonatanklosko
jonatanklosko merged commit 55ec9ac into elixir-nx:main Dec 18, 2025
2 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants