jax.profiler.stop_trace

jax.profiler.stop_trace()

停止当前正在运行的 profiler trace。

trace 将保存到对应 start_trace() 调用中传递的 log_dir。如果在启动 trace 之前调用，将引发 RuntimeError。