jax.profiler.save_device_memory_profile#
- jax.profiler.save_device_memory_profile(filename, backend=None)[source]#
收集设备内存配置文件并将其写入文件。
save_device_memory_profile()是围绕device_memory_profile()的一个便利包装器，它将其输出保存到filename中。有关更多信息，请参阅device_memory_profile()文档。- 参数：
filename – 应将配置文件写入到的文件名。
backend (str | None | None) – 可选；应为其收集设备内存配置文件的 JAX 后端的名称。
- 返回类型：
None