xla
3f3d01e7 - [Profiler] Introduce trace_me decorator (#4809)

Commit
2 years ago
[Profiler] Introduce trace_me decorator (#4809) Summary: This pull request introduces a new profiler API called trace_me decorator. I find it very tedious if I want to scope a method as I need to first use the xp.Trace("") context manager and then indent the whole method. Now, you can just do: ``` import torch_xla.debug.profiler as xp @xp.trace_me("forward") def forward: ... ``` Test Plan: PJRT_DEVICE=CPU python test/test_operations.py -v -k XpTraceTest
Author
Parents
Loading