Skip to content

fix(rfd3): cast symmetry frames to map dtype in batched collate - #419

Open
mayuriphad wants to merge 1 commit into
RosettaCommons:productionfrom
mayuriphad:fix/symmetry-translation-dtype
Open

mayuriphad wants to merge 1 commit into
RosettaCommons:productionfrom
mayuriphad:fix/symmetry-translation-dtype

Conversation

@mayuriphad

Copy link
Copy Markdown

Fixes #418

symmetry_maps in rfd3/inference/batching.py copies each frame's rotation and translation straight into the preallocated float32 maps. When a frame arrives as an integer tensor (for example the int64 zero translation in a C2 design), the index put fails with:

RuntimeError: Index put requires the source and destination dtypes match, got Float for the destination and Long for the source.

The legacy (non-batched) path does not hit this because it applies the frames with arithmetic, which promotes the dtype.

This casts r and t to the destination dtype and device before writing them.

Test

Added test_symmetry_maps_accept_integer_frames to models/rfd3/tests/test_rfd3_batch_symmetry.py: it collates an example whose frames are long tensors and checks the maps come back as float32 with the right translation. On production without the fix it fails with the error above; with the fix the full CPU suite passes (pytest, mypy and ruff green on my fork).

symmetry_maps wrote each frame's rotation and translation straight into
float32 tensors, so an int64 translation raised "Index put requires the
source and destination dtypes match" before the first diffusion step.

Fixes RosettaCommons#418

Signed-off-by: mayuriphad <mayuriphad656@gmail.com>
Copilot AI balanced review requested due to automatic review settings November 10, 2026 13:54

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota. Try requesting a review at a lower effort level.

An organization owner can choose the repository organization as the billing source for reviews requested by members with a paid Copilot license. Learn about Copilot code review billing.

This branch has not been deployed

No deployments
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.

[BUG] Symmetric designs fail in the batched path: the frame translation is int64

2 participants