Skip to content

fix: automatically fallback to chunk_size=1 when vmap fails#767

Open
KartavyaDikshit wants to merge 1 commit into
SimplexLab:mainfrom
KartavyaDikshit:fix-issue-699
Open

fix: automatically fallback to chunk_size=1 when vmap fails#767
KartavyaDikshit wants to merge 1 commit into
SimplexLab:mainfrom
KartavyaDikshit:fix-issue-699

Conversation

@KartavyaDikshit

Copy link
Copy Markdown

Resolves #699

Fallback to sequential differentiation when torch.vmap encounters limitations such as mixed precision autocast nodes in graph or torch.compile. Emits a helpful warning suggesting explicitly passing parallel_chunk_size=1.

Signed-off-by: KartavyaDikshit kartavyaniraj.dikshit2021@vitstudent.ac.in

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.

autojac incompatible with AMP for models with normalization layers

1 participant