named scopes for better profiling of memory #598
Description
Check for existing issues
- Completed
Description
For better profiling of device memory use jax.named_scope.
{{ message }}
For better profiling of device memory use jax.named_scope.