Skip to content

Cast weights to np.array before checkpointing - #51

Merged
jessegrabowski merged 2 commits into
mainfrom
mlx-checkpointing
Jul 30, 2026
Merged

Cast weights to np.array before checkpointing#51
jessegrabowski merged 2 commits into
mainfrom
mlx-checkpointing

Conversation

@jessegrabowski

Copy link
Copy Markdown
Member

Current our save/load checkpoint functions break on jax/mlx backend because we assume they will be numpy arrays. Easy enough to just cast them before writing.

@jessegrabowski jessegrabowski added bug Something isn't working enhancement New feature or request labels Jul 30, 2026
@jessegrabowski
jessegrabowski merged commit 10ccfb6 into main Jul 30, 2026
4 checks passed
@jessegrabowski
jessegrabowski deleted the mlx-checkpointing branch July 30, 2026 19:05
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

bug Something isn't working enhancement New feature or request

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant