Skip to content

How to do device memory profiling in Colab for JAX/Flax? #918

Answered by marcvanzee
marcvanzee asked this question in Q&A
How to do device memory profiling in Colab for JAX/Flax? #918
@marcvanzee marcvanzee
Jan 18, 2021 · 1 comments · 1 reply
Answered by marcvanzee Return to top
Discussion options

marcvanzee
Jan 18, 2021
Maintainer

You must be logged in to vote
Answered by marcvanzee Jan 21, 2021

Some simple tips:

  • JAX memory profiling can provide help to spot where the issue occurs.
  • you could run the cell with %%jaxheapprof to get the device memory profile from Colab.
View full answer

Replies: 1 comment 1 reply

Comment options

marcvanzee
Jan 21, 2021
Maintainer Author

You must be logged in to vote
1 reply
@ayaka14732
Comment options

ayaka14732 Feb 19, 2022
Collaborator

Answer selected by marcvanzee
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Category
Q&A
Labels
None yet
2 participants
@marcvanzee @ayaka14732
Converted from issue

This discussion was converted from issue #870 on January 21, 2021 14:04.