jax.profiler.save_device_memory_profile#
- jax.profiler.save_device_memory_profile(filename, backend=None)[source]#
Collects a device memory profile and writes it to a file.
save_device_memory_profile()
is a convenience wrapper arounddevice_memory_profile()
that saves its output to afilename
. See thedevice_memory_profile()
documentation for more information.